diff --git a/.buildkite/scripts/hardware_ci/run-xpu-test.sh b/.buildkite/scripts/hardware_ci/run-xpu-test.sh
index 2daf1534bd63..c1164bf43925 100644
--- a/.buildkite/scripts/hardware_ci/run-xpu-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-xpu-test.sh
@@ -44,7 +44,7 @@ docker run \
python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2
python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
cd tests
- pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py
+ pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py
pytest -v -s v1/engine
pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py
diff --git a/.buildkite/scripts/upload-release-wheels-pypi.sh b/.buildkite/scripts/upload-release-wheels-pypi.sh
index dacdb6e92483..058e5bbe4f4c 100644
--- a/.buildkite/scripts/upload-release-wheels-pypi.sh
+++ b/.buildkite/scripts/upload-release-wheels-pypi.sh
@@ -54,10 +54,13 @@ mkdir -p $DIST_DIR
# include only wheels for the release version, ignore all files with "dev" or "rc" in the name (without excluding 'aarch64')
aws s3 cp --recursive --exclude "*" --include "vllm-${PURE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc[0-9]*" "$S3_COMMIT_PREFIX" $DIST_DIR
echo "Wheels copied to local directory"
-# generate source tarball
-git archive --format=tar.gz --output="$DIST_DIR/vllm-${PURE_VERSION}.tar.gz" "$BUILDKITE_COMMIT"
+# generate source distribution using setup.py
+python setup.py sdist --dist-dir=$DIST_DIR
ls -la $DIST_DIR
+SDIST_FILE=$(find $DIST_DIR -name "vllm*.tar.gz")
+echo "Found sdist: $SDIST_FILE"
+
# upload wheels to PyPI (only default variant, i.e. files without '+' in the name)
PYPI_WHEEL_FILES=$(find $DIST_DIR -name "vllm-${PURE_VERSION}*.whl" -not -name "*+*")
if [[ -z "$PYPI_WHEEL_FILES" ]]; then
@@ -65,6 +68,6 @@ if [[ -z "$PYPI_WHEEL_FILES" ]]; then
exit 1
fi
-python3 -m twine check "$PYPI_WHEEL_FILES"
-python3 -m twine upload --non-interactive --verbose "$PYPI_WHEEL_FILES"
-echo "Wheels uploaded to PyPI"
+python3 -m twine check "$PYPI_WHEEL_FILES" "$SDIST_FILE"
+python3 -m twine upload --non-interactive --verbose "$PYPI_WHEEL_FILES" "$SDIST_FILE"
+echo "Wheels and source distribution uploaded to PyPI"
diff --git a/requirements/common.txt b/requirements/common.txt
index 9ee1b71512b1..a8e1450fa1ee 100644
--- a/requirements/common.txt
+++ b/requirements/common.txt
@@ -12,7 +12,7 @@ tokenizers >= 0.21.1 # Required for fast incremental detokenization.
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.
aiohttp >= 3.13.3
-openai >= 1.99.1 # For Responses API with reasoning content
+openai >= 1.99.1, < 2.25.0 # For Responses API with reasoning content
pydantic >= 2.12.0
prometheus_client >= 0.18.0
pillow # Required for image processing
diff --git a/requirements/cuda.txt b/requirements/cuda.txt
index 22477dc821b7..2dd380d148c4 100644
--- a/requirements/cuda.txt
+++ b/requirements/cuda.txt
@@ -11,6 +11,9 @@ torchaudio==2.10.0
torchvision==0.25.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
# FlashInfer should be updated together with the Dockerfile
flashinfer-python==0.6.4
+# Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to
+# breaking changes in 1.19.0
+nvidia-cudnn-frontend>=1.13.0,<1.19.0
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
nvidia-cutlass-dsl>=4.4.0.dev1
diff --git a/requirements/xpu.txt b/requirements/xpu.txt
index 050737164f8c..3271f9f39275 100644
--- a/requirements/xpu.txt
+++ b/requirements/xpu.txt
@@ -15,4 +15,4 @@ torch==2.10.0+xpu
torchaudio
torchvision
-vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.2/vllm_xpu_kernels-0.1.2-cp312-cp312-linux_x86_64.whl
+vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.3/vllm_xpu_kernels-0.1.3-cp38-abi3-linux_x86_64.whl
diff --git a/tests/reasoning/test_nemotron_v3_reasoning_parser.py b/tests/reasoning/test_nemotron_v3_reasoning_parser.py
new file mode 100644
index 000000000000..c7ba95cb11bd
--- /dev/null
+++ b/tests/reasoning/test_nemotron_v3_reasoning_parser.py
@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import TypedDict
+
+import pytest
+import regex as re
+
+from tests.reasoning.utils import run_reasoning_extraction
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.reasoning import ReasoningParser, ReasoningParserManager
+
+parser_name = "nemotron_v3"
+
+
+class ReasoningCase(TypedDict):
+ output: str
+ reasoning: str | None
+ content: str | None
+
+
+class FakeNemotronTokenizer:
+ def __init__(self):
+ self._vocab = {
+ "": 1,
+ "": 2,
+ }
+ self._pattern = re.compile(r"(|)")
+
+ def get_vocab(self) -> dict[str, int]:
+ return self._vocab
+
+ def tokenize(self, text: str) -> list[str]:
+ tokens: list[str] = []
+ for part in self._pattern.split(text):
+ if part:
+ tokens.append(part)
+ return tokens
+
+ def convert_tokens_to_string(self, tokens: list[str]) -> str:
+ return "".join(tokens)
+
+
+@pytest.fixture
+def tokenizer():
+ return FakeNemotronTokenizer()
+
+
+@pytest.mark.parametrize(
+ "streaming,param_dict",
+ [
+ pytest.param(
+ False,
+ {
+ "output": "This is a reasoning sectionThis is the rest",
+ "reasoning": "This is a reasoning section",
+ "content": "This is the rest",
+ },
+ id="without_start_token",
+ ),
+ pytest.param(
+ True,
+ {
+ "output": "This is a reasoning sectionThis is the rest",
+ "reasoning": "This is a reasoning section",
+ "content": "This is the rest",
+ },
+ id="without_start_token_streaming",
+ ),
+ pytest.param(
+ False,
+ {
+ "output": "This is a reasoning sectionThis is the rest",
+ "reasoning": "This is a reasoning section",
+ "content": "This is the rest",
+ },
+ id="with_start_token",
+ ),
+ pytest.param(
+ True,
+ {
+ "output": "This is a reasoning sectionThis is the rest",
+ "reasoning": "This is a reasoning section",
+ "content": "This is the rest",
+ },
+ id="with_start_token_streaming",
+ ),
+ ],
+)
+def test_nemotron_v3_reasoning(
+ tokenizer: FakeNemotronTokenizer,
+ streaming: bool,
+ param_dict: ReasoningCase,
+):
+ output = tokenizer.tokenize(param_dict["output"])
+ model_output = [tokenizer.convert_tokens_to_string([token]) for token in output]
+ parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
+ tokenizer
+ )
+
+ reasoning, content = run_reasoning_extraction(
+ parser, model_output, streaming=streaming
+ )
+
+ assert reasoning == param_dict["reasoning"]
+ assert content == param_dict["content"]
+
+
+def test_nemotron_v3_without_thinking_returns_content(
+ tokenizer: FakeNemotronTokenizer,
+):
+ parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
+ parser = parser_cls(tokenizer)
+ request = ChatCompletionRequest(
+ model="test-model",
+ messages=[],
+ chat_template_kwargs={"enable_thinking": False},
+ )
+
+ reasoning, content = run_reasoning_extraction(
+ parser,
+ ["This is plain content"],
+ request=request,
+ streaming=False,
+ )
+
+ assert reasoning is None
+ assert content == "This is plain content"
+
+
+def test_nemotron_v3_force_nonempty_content_returns_content(
+ tokenizer: FakeNemotronTokenizer,
+):
+ parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
+ parser = parser_cls(tokenizer)
+ request = ChatCompletionRequest(
+ model="test-model",
+ messages=[],
+ chat_template_kwargs={"force_nonempty_content": True},
+ )
+
+ reasoning, content = run_reasoning_extraction(
+ parser,
+ ["This is plain content"],
+ request=request,
+ streaming=False,
+ )
+
+ assert reasoning is None
+ assert content == "This is plain content"
+
+
+def test_nemotron_v3_with_thinking_keeps_truncated_reasoning(
+ tokenizer: FakeNemotronTokenizer,
+):
+ parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
+ parser = parser_cls(tokenizer)
+ request = ChatCompletionRequest(
+ model="test-model",
+ messages=[],
+ chat_template_kwargs={"enable_thinking": True},
+ )
+
+ reasoning, content = run_reasoning_extraction(
+ parser,
+ ["This is truncated reasoning"],
+ request=request,
+ streaming=False,
+ )
+
+ assert reasoning == "This is truncated reasoning"
+ assert content is None
diff --git a/tests/tokenizers_/test_basic.py b/tests/tokenizers_/test_basic.py
index 1c1dd33381f9..cf0d8f53c6f2 100644
--- a/tests/tokenizers_/test_basic.py
+++ b/tests/tokenizers_/test_basic.py
@@ -29,7 +29,8 @@ def test_tokenizer_like_protocol():
_assert_tokenizer_like(tokenizer)
tokenizer = get_tokenizer(
- "mistralai/Mistral-7B-Instruct-v0.3", tokenizer_mode="mistral"
+ "mistralai/Mistral-7B-Instruct-v0.3",
+ tokenizer_mode="mistral",
)
assert isinstance(tokenizer, MistralTokenizer)
_assert_tokenizer_like(tokenizer)
@@ -40,11 +41,20 @@ def test_tokenizer_like_protocol():
tokenizer = get_tokenizer("deepseek-ai/DeepSeek-V3", tokenizer_mode="deepseek_v32")
assert isinstance(tokenizer, HfTokenizer)
+
# Verify it's a fast tokenizer (required for FastIncrementalDetokenizer)
assert isinstance(tokenizer, PreTrainedTokenizerFast)
assert "DSV32" in tokenizer.__class__.__name__
_assert_tokenizer_like(tokenizer)
+ tokenizer = get_tokenizer(
+ "Qwen/Qwen-VL",
+ tokenizer_mode="qwen_vl",
+ trust_remote_code=True,
+ )
+ assert isinstance(tokenizer, HfTokenizer)
+ assert "WithoutImagePad" in tokenizer.__class__.__name__
+
@pytest.mark.parametrize("tokenizer_name", ["facebook/opt-125m", "gpt2"])
def test_tokenizer_revision(tokenizer_name: str):
diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py
index f8bf52de0832..7c9a95ef1262 100644
--- a/vllm/benchmarks/serve.py
+++ b/vllm/benchmarks/serve.py
@@ -1321,6 +1321,7 @@ def add_cli_args(parser: argparse.ArgumentParser):
- "slow" will always use the slow tokenizer.\n
- "mistral" will always use the tokenizer from `mistral_common`.\n
- "deepseek_v32" will always use the tokenizer from `deepseek_v32`.\n
+ - "qwen_vl" will always use the tokenizer from `qwen_vl`.\n
- Other custom values can be supported via plugins.""",
)
parser.add_argument("--use-beam-search", action="store_true")
diff --git a/vllm/config/model.py b/vllm/config/model.py
index 4e3568fa15b1..6c48bfde6437 100644
--- a/vllm/config/model.py
+++ b/vllm/config/model.py
@@ -126,6 +126,7 @@ class ModelConfig:
- "slow" will always use the slow tokenizer.\n
- "mistral" will always use the tokenizer from `mistral_common`.\n
- "deepseek_v32" will always use the tokenizer from `deepseek_v32`.\n
+ - "qwen_vl" will always use the tokenizer from `qwen_vl`.\n
- Other custom values can be supported via plugins."""
trust_remote_code: bool = False
"""Trust remote code (e.g., from HuggingFace) when downloading the model
diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py
index a950ba531ad2..1eea44f9f288 100644
--- a/vllm/config/speculative.py
+++ b/vllm/config/speculative.py
@@ -299,16 +299,16 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
{"n_predict": n_predict, "architectures": ["ExaoneMoeMTP"]}
)
- if hf_config.model_type in ("qwen3_5", "qwen3_5_moe"):
- is_moe = hf_config.model_type == "qwen3_5_moe"
- hf_config.model_type = "qwen3_5_mtp"
- n_predict = getattr(hf_config, "mtp_num_hidden_layers", None)
- hf_config.update(
- {
- "n_predict": n_predict,
- "architectures": ["Qwen3_5MoeMTP" if is_moe else "Qwen3_5MTP"],
- }
- )
+ # if hf_config.model_type in ("qwen3_5", "qwen3_5_moe"):
+ # is_moe = hf_config.model_type == "qwen3_5_moe"
+ # hf_config.model_type = "qwen3_5_mtp"
+ # n_predict = getattr(hf_config, "mtp_num_hidden_layers", None)
+ # hf_config.update(
+ # {
+ # "n_predict": n_predict,
+ # "architectures": ["Qwen3_5MoeMTP" if is_moe else "Qwen3_5MTP"],
+ # }
+ # )
if hf_config.model_type == "longcat_flash":
hf_config.model_type = "longcat_flash_mtp"
n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)
@@ -347,7 +347,7 @@ def __post_init__(self):
"method `%s` is deprecated and replaced with mtp.", self.method
)
self.method = "mtp"
-
+
if self.model is None and self.num_speculative_tokens is not None:
if self.method == "mtp":
if self.target_model_config is None:
@@ -465,7 +465,7 @@ def __post_init__(self):
)
# Automatically detect the method
- if self.method in ("eagle", "eagle3"):
+ if self.method in ("eagle", "eagle3", 'medusa'):
pass
# examples:
# yuhuili/EAGLE-LLaMA3-Instruct-8B
@@ -589,6 +589,7 @@ def __post_init__(self):
self.target_parallel_config, self.draft_tensor_parallel_size
)
)
+ self.method = "medusa"
return self
def _validate_suffix_decoding(self):
diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
index febb3b2ef0d7..1ed76f8920cf 100644
--- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
@@ -35,12 +35,6 @@ def __init__(
):
super().__init__(moe_config, quant_config)
- if moe_config.moe_parallel_config.use_ep and quant_config.is_per_tensor:
- raise NotImplementedError(
- "EP parallelism is not supported with TRTLLM"
- "per-tensor FP8 quantization."
- )
-
self.routing_method_type = moe_config.routing_method
self.topk = moe_config.experts_per_token
self.intermediate_size_per_partition = (
@@ -182,9 +176,6 @@ def _apply_per_block(
assert not apply_router_weight_on_input
assert activation == MoEActivation.SILU
- if e_score_correction_bias is not None:
- e_score_correction_bias = e_score_correction_bias.to(hidden_states.dtype)
-
if self.routing_method_type == RoutingMethodType.DeepSeekV3:
router_logits = router_logits.to(torch.float32)
@@ -240,12 +231,11 @@ def _apply_per_tensor(
) -> torch.Tensor:
# Delay import for non-CUDA.
import flashinfer
- from flashinfer.fused_moe.core import ActivationType
# Confirm supported activation function.
assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
- activation_type = ActivationType(activation_to_flashinfer_int(activation))
+ activation_type = activation_to_flashinfer_int(activation)
# Confirm Llama-4 routing is proper.
if self.routing_method_type == RoutingMethodType.Llama4:
diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
index 502671766400..174c581b396f 100644
--- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
@@ -323,4 +323,5 @@ def apply(
routed_scaling_factor=routed_scaling_factor,
routing_method_type=self.routing_method_type,
do_finalize=True,
+ activation_type=activation_to_flashinfer_int(activation),
)[0]
diff --git a/vllm/model_executor/layers/fused_moe/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/fused_batched_moe.py
index 68393f768dcc..b6441552a4e1 100644
--- a/vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_batched_moe.py
@@ -912,7 +912,7 @@ def _supports_current_device() -> bool:
@staticmethod
def _supports_no_act_and_mul() -> bool:
- return False
+ return True
@staticmethod
def _supports_quant_scheme(
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
index 023cdd0b4340..ee321f241aad 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -1944,7 +1944,7 @@ def _supports_current_device() -> bool:
@staticmethod
def _supports_no_act_and_mul() -> bool:
- return False
+ return True
@staticmethod
def _supports_quant_scheme(
@@ -1983,6 +1983,9 @@ def _supports_activation(activation: MoEActivation) -> bool:
MoEActivation.GELU,
MoEActivation.SWIGLUOAI,
MoEActivation.SWIGLUSTEP,
+ MoEActivation.SILU_NO_MUL,
+ MoEActivation.GELU_NO_MUL,
+ MoEActivation.RELU2_NO_MUL,
]
@staticmethod
diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py
index a696d2129c28..4560237250aa 100644
--- a/vllm/model_executor/models/internvl.py
+++ b/vllm/model_executor/models/internvl.py
@@ -17,7 +17,7 @@
import torchvision.transforms as T
from PIL import Image
from transformers import BatchFeature, PretrainedConfig, TensorType
-
+import os
from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -60,6 +60,8 @@
)
from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
+MAX_PATCHES_WITH_RESIZE = os.environ.get("MAX_PATCHES_WITH_RESIZE", "0").lower() in ("1", "true", "yes")
+print("MAX_PATCHES_WITH_RESIZE: {}".format(MAX_PATCHES_WITH_RESIZE))
IMG_START = "![]()
"
IMG_END = ""
IMG_CONTEXT = ""
@@ -167,6 +169,46 @@ def find_closest_aspect_ratio(
best_ratio = ratio
return best_ratio
+def _find_closest_aspect_ratio_optim(aspect_ratio, target_ratios, width, height, image_size,
+ top_k=3, ar_threshold=0.2):
+ """
+ 先按面积接近程度选 top_k 个候选，再从中选长宽比最接近的。
+
+ Args:
+ ar_threshold: 可选，长宽比差异阈值。若设置，则只有与原图长宽比差异
+ 不超过该阈值的 ratio 才能进入候选列表。
+ """
+ area = width * height
+ # 1. 按面积差异排序，可选地过滤长宽比差异过大的 ratio
+ candidates = []
+ for ratio in target_ratios:
+ target_ar = ratio[0] / ratio[1]
+ ar_diff = abs(aspect_ratio - target_ar)
+
+ if ar_threshold is not None and ar_diff > ar_threshold:
+ continue
+
+ target_area = image_size * image_size * ratio[0] * ratio[1]
+ area_diff = abs(area - target_area)
+ candidates.append((ratio, area_diff, ar_diff))
+
+ # 如果阈值过滤后没有候选，回退到不过滤
+ if not candidates:
+ for ratio in target_ratios:
+ target_ar = ratio[0] / ratio[1]
+ ar_diff = abs(aspect_ratio - target_ar)
+ target_area = image_size * image_size * ratio[0] * ratio[1]
+ area_diff = abs(area - target_area)
+ candidates.append((ratio, area_diff, ar_diff))
+
+ # 2. 按面积差异排序，取 top_k
+ candidates.sort(key=lambda x: x[1])
+ top_candidates = candidates[:top_k]
+
+ # 3. 从 top_k 中选长宽比最接近的
+ top_candidates.sort(key=lambda x: x[2])
+ return top_candidates[0][0]#, top_candidates[0][-1]
+
def resolve_internvl_min_max_num(
*,
@@ -209,19 +251,29 @@ def calculate_internvl_targets(
aspect_ratio = orig_width / orig_height
# find the closest aspect ratio to the target
- target_aspect_ratio = find_closest_aspect_ratio(
+
+ # print(MAX_PATCHES_WITH_RESIZE)
+ if MAX_PATCHES_WITH_RESIZE:
+ target_aspect_ratio = _find_closest_aspect_ratio_optim(
aspect_ratio,
target_ratios,
width=orig_width,
height=orig_height,
image_size=image_size,
)
+ else:
+ target_aspect_ratio = find_closest_aspect_ratio(
+ aspect_ratio,
+ target_ratios,
+ width=orig_width,
+ height=orig_height,
+ image_size=image_size,
+ )
# calculate the target width and height
target_width = image_size * target_aspect_ratio[0]
target_height = image_size * target_aspect_ratio[1]
blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
-
# add thumbnail image if num_blocks != 1
if use_thumbnail and blocks != 1:
blocks += 1
@@ -249,6 +301,7 @@ def dynamic_preprocess_internvl(
)
# resize the image
+
resized_img = image.resize((target_width, target_height))
processed_images = []
for i in range(blocks):
@@ -262,11 +315,14 @@ def dynamic_preprocess_internvl(
split_img = resized_img.crop(box)
processed_images.append(split_img)
+
assert len(processed_images) == blocks
if use_thumbnail and len(processed_images) != 1:
thumbnail_img = image.resize((image_size, image_size))
processed_images.append(thumbnail_img)
+ #for i in range(len(processed_images)):
+ # processed_images[i].save("tmp/debug_tile_{}.png".format(i))
return processed_images
@@ -362,6 +418,7 @@ def __init__(
self.num_image_token = int(
(image_size // patch_size) ** 2 * (config.downsample_ratio**2)
)
+
self.image_size = image_size
self.min_dynamic_patch = min_dynamic_patch
self.max_dynamic_patch = max_dynamic_patch
@@ -482,6 +539,7 @@ def _preprocess_image(
if len(images) == 0:
image_inputs = {}
else:
+
pixel_values_lst = self._images_to_pixel_values_lst(
images,
min_dynamic_patch=min_dynamic_patch,
@@ -494,11 +552,13 @@ def _preprocess_image(
[len(item) for item in pixel_values_lst]
),
}
+ #print(max_dynamic_patch, min_dynamic_patch)
+ #print(images[0].size, image_inputs['pixel_values_flat'].shape, image_inputs['image_num_patches']) # torch.Size([49, 3, 448, 448]) tensor([49])
for pixel_values in pixel_values_lst:
num_patches = pixel_values.shape[0]
feature_size = num_patches * self.num_image_token
-
+
image_repl = self.get_image_repl(feature_size, num_patches)
text = [t.replace("", image_repl.full, 1) for t in text]
return text, image_inputs
@@ -520,7 +580,6 @@ def __call__(
return_tensors: str | TensorType | None = None,
) -> BatchFeature:
text, images = [self._make_batch_input(x) for x in (text, images)]
-
text, image_inputs = self._preprocess_image(
text=text,
images=images,
diff --git a/vllm/model_executor/models/medusa.py b/vllm/model_executor/models/medusa.py
index fd7fc2c73f16..1f4ba9abf7a0 100644
--- a/vllm/model_executor/models/medusa.py
+++ b/vllm/model_executor/models/medusa.py
@@ -12,168 +12,508 @@
ParallelLMHead,
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.qwen3_5 import *
+from vllm.model_executor.models.qwen3_5_mtp import *
+from vllm.model_executor.models.qwen3_next import *
+from vllm.model_executor.models.qwen3 import *
+
from .utils import maybe_prefix
-class ResidualBlock(nn.Module):
- def __init__(self, config: VllmConfig, hidden_size: int, num_layers: int) -> None:
- super().__init__()
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-only Qwen3_5 MTP model."""
- self.layers = nn.ModuleList(
- [
- nn.Linear(
- hidden_size,
- hidden_size,
- bias=getattr(config, "medusa_fc_bias", False),
- )
- for _ in range(num_layers)
- ]
- )
- self.act = nn.SiLU()
-
- def forward(self, x: torch.Tensor) -> torch.Tensor:
- for layer in self.layers:
- x = x + self.act(layer(x))
- return x
-
-
-class Medusa(nn.Module):
- """This class implements the Medusa draft model from the paper: https://arxiv.org/abs/2401.10774
- Reference implementation: https://github.com/FasterDecoding/Medusa
-
- Differences from reference implementation:
- 1. Currently this only supports generating proposals from top-1 tokens.
- 2. We have an optional token_map which reduces draft vocab to most
- frequently used tokens to give some additional speed-up by reducing
- sampling overhead. This is disabled unless the checkpoint file has
- explicit token_map tensor and config has an optional attribute
- truncated_vocab_size < vocab_size. To use this technique, one has to find
- the top-k most frequent tokens in target dataset and add that as a tensor
- in the draft checkpoint (using key token_map). Also, the draft config
- needs to have truncated_vocab_size (=k) as an attribute."""
-
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
- config = vllm_config.speculative_config.draft_model_config.hf_config
+import typing
+from collections.abc import Callable, Iterable
+
+import torch
+from torch import nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.layers.linear import ColumnParallelLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+ ParallelLMHead,
+ VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.qwen3 import RMSNorm
+from vllm.sequence import IntermediateTensors
+
+from .interfaces import (
+ MultiModalEmbeddings,
+ SupportsMultiModal,
+ _require_is_multimodal,
+)
+from .utils import (
+ AutoWeightsLoader,
+ PPMissingLayer,
+ _merge_multimodal_embeddings,
+ is_pp_missing_parameter,
+ make_empty_intermediate_tensors_factory,
+ maybe_prefix,
+)
+
+logger = init_logger(__name__)
+
+
+class MTPMLPModule(nn.Module):
+ def __init__(self, config, quant_config, prefix):
+ from vllm.model_executor.models.qwen3 import Qwen3MLP
super().__init__()
self.config = config
- self.blocks = nn.ModuleList(
- [
- ResidualBlock(
- config=config,
- hidden_size=self.config.hidden_size,
- num_layers=self.config.num_hidden_layers,
+ self.mlp = Qwen3MLP(
+ hidden_size=config.hidden_size,
+ intermediate_size=config.intermediate_size,
+ hidden_act=config.hidden_act,
+ quant_config=quant_config,
+ prefix=f"{prefix}.mlp",
)
- for _ in range(self.config.num_heads)
- ]
+
+ def forward(self,
+ hidden_states: torch.Tensor,
+ input_ids: torch.Tensor | None = None,
+ positions: torch.Tensor | None = None,
+ intermediate_tensors: IntermediateTensors | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ **kwargs: object,
+ ):
+ return self.mlp(hidden_states), hidden_states
+
+
+class Qwen3_5MultiTokenPredictor(nn.Module):
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+ super().__init__()
+
+ model_config = vllm_config.model_config
+ quant_config = vllm_config.quant_config
+
+ config = model_config.hf_text_config
+
+ self.config = config
+
+ self.vocab_size = config.vocab_size
+
+ self.mtp_start_layer_idx = config.num_hidden_layers
+ self.num_mtp_layers = getattr(config, "mtp_num_hidden_layers", 1)
+
+ self.embed_tokens = VocabParallelEmbedding(
+ self.vocab_size,
+ config.hidden_size,
)
- self.orig_vocab_size = config.vocab_size
- self.truncated_vocab_size = config.truncated_vocab_size
-
- if getattr(config, "original_lm_head", False):
- self.lm_head = ParallelLMHead(
- self.truncated_vocab_size,
- config.hidden_size,
- prefix=maybe_prefix(prefix, "lm_head"),
- )
- self.lm_heads = [self.lm_head for _ in range(self.config.num_heads)]
- else:
- self.lm_heads = nn.ModuleList(
- [
- ParallelLMHead(
- config.vocab_size,
- config.hidden_size,
- prefix=maybe_prefix(prefix, f"lm_heads.{i}"),
- )
- for i in range(self.config.num_heads)
- ]
+
+ self.fc = ColumnParallelLinear(
+ self.config.hidden_size * 2,
+ self.config.hidden_size,
+ gather_output=True,
+ bias=False,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.fc",
+ )
+
+ self.layers = torch.nn.ModuleList(
+ MTPMLPModule(
+ config=config,
+ quant_config=quant_config,
+ prefix=f"{prefix}.layers.{idx}",
)
+ for idx in range(self.num_mtp_layers)
+ )
- logit_scale = getattr(config, "logit_scale", 1.0)
- self.logits_processor = LogitsProcessor(
- config.vocab_size, self.truncated_vocab_size, logit_scale
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+ ["hidden_states", "residual"], config.hidden_size
)
- # Token map is a idx to token mapping to reduce the vocab size for
- # the draft model. Using smaller vocab size for draft, containing
- # only most frequent tokens reduces the speculation overhead. This
- # doesn't affect the acceptance rate much and thus gives more speed
- # -up. By default, this is disabled and is only used if the EAGLE
- # checkpoint file has token_map tensor.
- self.token_map = None
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.pre_fc_norm_hidden = RMSNorm(
+ config.hidden_size, eps=config.rms_norm_eps
+ )
+ self.pre_fc_norm_embedding = RMSNorm(
+ config.hidden_size, eps=config.rms_norm_eps
+ )
- def forward(self, hidden_states: torch.Tensor) -> list[torch.Tensor]:
- return [block(hidden_states) for block in self.blocks]
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+ return self.embed_tokens(input_ids)
- def compute_logits(
+ def forward(
self,
- hidden_states: list[torch.Tensor],
- ) -> list[torch.Tensor]:
- logits_lst: list[torch.Tensor] = []
+ hidden_states: torch.Tensor,
+ inputs_embeds: torch.Tensor | None = None,
+ input_ids: torch.Tensor | None = None,
+ positions: torch.Tensor | None = None,
+ intermediate_tensors: IntermediateTensors | None = None,
+ spec_step_idx: int = 0,
+ ) -> torch.Tensor:
+ if get_pp_group().is_first_rank:
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_input_ids(input_ids)
+ assert hidden_states.shape[-1] == inputs_embeds.shape[-1]
+ inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds)
+ hidden_states = self.pre_fc_norm_hidden(hidden_states)
+ # hidden_states = torch.cat([inputs_embeds, hidden_states], dim=-1)
+ hidden_states = torch.cat([hidden_states, inputs_embeds], dim=-1)
+ hidden_states = self.fc(hidden_states)
+ residual = None
+ else:
+ assert intermediate_tensors is not None
+ hidden_states = intermediate_tensors["hidden_states"]
+ residual = intermediate_tensors["residual"]
- for hs, lm_head in zip(hidden_states, self.lm_heads):
- _logits = self.logits_processor(lm_head, hs)
+ current_step_idx = spec_step_idx % self.num_mtp_layers
+ hidden_states, residual = self.layers[current_step_idx](
+ positions=positions,
+ hidden_states=hidden_states,
+ residual=residual,
+ )
- if _logits is None:
- # _logits should only be None on rank > 0, in which case
- # it should remain true for every lm_head
- assert len(logits_lst) == 0
- continue
+ if not get_pp_group().is_last_rank:
+ return IntermediateTensors(
+ {"hidden_states": hidden_states, "residual": residual}
+ )
- if self.token_map is None:
- logits_lst.append(_logits)
- else:
- logits_lst.append(
- -torch.inf
- * torch.ones(
- size=(*_logits.shape[:-1], self.orig_vocab_size),
- device=_logits.device,
- dtype=_logits.dtype,
- )
- )
+ hidden_states, _ = self.norm(hidden_states, residual)
+ return hidden_states
- logits_lst[-1][..., self.token_map] = _logits
+ def load_fused_expert_weights(
+ self,
+ name: str,
+ params_dict: dict,
+ loaded_weight: torch.Tensor,
+ shard_id: str,
+ num_experts: int,
+ ) -> bool:
+ param = params_dict[name]
+ weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
+ loaded_local_expert = False
+ for expert_id in range(num_experts):
+ curr_expert_weight = loaded_weight[expert_id]
+ success = weight_loader(
+ param,
+ curr_expert_weight,
+ name,
+ shard_id,
+ expert_id,
+ return_success=True,
+ )
+ if success:
+ loaded_local_expert = True
- return logits_lst
+ return loaded_local_expert
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ stacked_params_mapping = [
+ # (param_name, shard_name, shard_id)
+ ("qkv_proj", "q_proj", "q"),
+ ("qkv_proj", "k_proj", "k"),
+ ("qkv_proj", "v_proj", "v"),
+ ("gate_up_proj", "gate_proj", 0),
+ ("gate_up_proj", "up_proj", 1),
+ ]
+
+ # Params for weights, fp8 weight scales, fp8 activation scales
+ # (param_name, weight_name, expert_id, shard_id)
+ expert_params_mapping = FusedMoE.make_expert_params_mapping(
+ self,
+ ckpt_gate_proj_name="gate_proj",
+ ckpt_down_proj_name="down_proj",
+ ckpt_up_proj_name="up_proj",
+ num_experts=self.config.num_experts
+ if hasattr(self.config, "num_experts")
+ else 0,
+ )
+
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
+ is_fused_expert = False
+ fused_expert_params_mapping = [
+ ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"),
+ ("experts.w2_weight", "experts.down_proj", 0, "w2"),
+ ]
+ num_experts = (
+ self.config.num_experts if hasattr(self.config, "num_experts") else 0
+ )
+ for name, loaded_weight in weights:
+ if "rotary_emb.inv_freq" in name:
+ continue
- weights_map = {}
+ for param_name, weight_name, shard_id in stacked_params_mapping:
+ if "experts.gate_up_proj" in name or "experts.down_proj" in name:
+ is_fused_expert = True
+ expert_params_mapping = fused_expert_params_mapping
- for name, loaded_weight in weights:
- name = name.replace("medusa_heads.", "")
-
- if name == "token_map":
- if self.truncated_vocab_size < self.orig_vocab_size:
- self.token_map = nn.Parameter(loaded_weight, requires_grad=False)
- elif name in params_dict:
- weights_map[name] = loaded_weight
- elif (
- getattr(self.config, "original_lm_head", False)
- and name == "lm_heads.0.weight"
- ):
- weights_map["lm_head.weight"] = loaded_weight
-
- for name, loaded_weight in weights_map.items():
- if (
- "lm_head" in name
- and self.token_map is not None
- and loaded_weight.shape[0] > self.token_map.shape[0]
- ):
- loaded_weight = loaded_weight[self.token_map]
-
- param = params_dict[name]
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
- weight_loader(param, loaded_weight)
+ if weight_name not in name:
+ continue
+
+ if "mlp.experts" in name:
+ continue
+
+ name = name.replace(weight_name, param_name)
+ # Skip loading extra bias for GPTQ models.
+ if name.endswith(".bias") and name not in params_dict:
+ continue
+ # Skip layers on other devices.
+ if is_pp_missing_parameter(name, self):
+ continue
+ if name not in params_dict:
+ continue
+ param = params_dict[name]
+ weight_loader = param.weight_loader
+ weight_loader(param, loaded_weight, shard_id)
+ break
+ else:
+ is_expert_weight = False
+ for mapping in expert_params_mapping:
+ param_name, weight_name, expert_id, shard_id = mapping
+ if weight_name not in name:
+ continue
+ is_expert_weight = True
+ name_mapped = name.replace(weight_name, param_name)
+ # Skip layers on other devices.
+ if is_pp_missing_parameter(name_mapped, self):
+ continue
+ if is_fused_expert:
+ # qwen3.5 no need to transpose
+ # loaded_weight = loaded_weight.transpose(-1, -2)
+ if "experts.gate_up_proj" in name:
+ loaded_weight = loaded_weight.chunk(2, dim=-2)
+ success_w1 = self.load_fused_expert_weights(
+ name_mapped,
+ params_dict,
+ loaded_weight[0],
+ "w1",
+ num_experts,
+ )
+ success_w3 = self.load_fused_expert_weights(
+ name_mapped,
+ params_dict,
+ loaded_weight[1],
+ "w3",
+ num_experts,
+ )
+ success = success_w1 and success_w3
+ else:
+ # down_proj
+ success = self.load_fused_expert_weights(
+ name_mapped,
+ params_dict,
+ loaded_weight,
+ shard_id,
+ num_experts,
+ )
+ if success:
+ name = name_mapped
+ break
+ else:
+ # Skip loading extra bias for GPTQ models.
+ if (
+ name_mapped.endswith(".bias")
+ or name_mapped.endswith("_bias")
+ ) and name_mapped not in params_dict:
+ continue
+ param = params_dict[name_mapped]
+ weight_loader = param.weight_loader
+ success = weight_loader(
+ param,
+ loaded_weight,
+ name_mapped,
+ shard_id=shard_id,
+ expert_id=expert_id,
+ return_success=True,
+ )
+ if success:
+ name = name_mapped
+ break
+ else:
+ if is_expert_weight:
+ # We've checked that this is an expert weight
+ # However it's not mapped locally to this rank
+ # So we simply skip it
+ continue
+ # Skip loading extra bias for GPTQ models.
+ if name.endswith(".bias") and name not in params_dict:
+ continue
+ if is_pp_missing_parameter(name, self):
+ continue
+ if name not in params_dict:
+ logger.warning_once(
+ f"Parameter {name} not found in params_dict, skip loading"
+ )
+ continue
+ param = params_dict[name]
+ weight_loader = getattr(
+ param, "weight_loader", default_weight_loader
+ )
+ weight_loader(param, loaded_weight)
loaded_params.add(name)
+ return loaded_params
- if self.token_map is not None:
- self.token_map.to(device=self.lm_heads[0].weight.device)
- assert (self.truncated_vocab_size == self.orig_vocab_size) or (
- self.token_map is not None
+class Medusa(nn.Module, SupportsMultiModal):
+ packed_modules_mapping = {
+ "qkv_proj": [
+ "q_proj",
+ "k_proj",
+ "v_proj",
+ ],
+ "gate_up_proj": ["gate_proj", "up_proj"],
+ }
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+ config = vllm_config.model_config.hf_text_config
+ self.vllm_config = vllm_config
+ cache_config = vllm_config.cache_config
+ if cache_config.mamba_cache_mode == "all":
+ raise NotImplementedError(
+ "Medusa currently does not support 'all' prefix caching, "
+ "please use '--mamba-cache-mode=align' instead"
+ )
+
+ self.quant_config = vllm_config.quant_config
+
+ super().__init__()
+ self.config = config
+ print('weishu debug medusa model load vllm_config', vllm_config)
+ self.num_mtp_tokens = vllm_config.speculative_config.num_speculative_tokens
+ self.model = Qwen3_5MultiTokenPredictor(
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "mtp")
)
+ print(f"[MEDUSA DEBUG] Medusa init: config.vocab_size={config.vocab_size}, "
+ f"model.vocab_size={self.model.vocab_size}, "
+ f"embed_tokens.num_embeddings={self.model.embed_tokens.num_embeddings}, "
+ f"num_mtp_tokens={self.num_mtp_tokens}")
- return loaded_params
+ if get_pp_group().is_last_rank:
+ if config.tie_word_embeddings:
+ self.lm_head = self.model.embed_tokens
+ else:
+ self.lm_head = ParallelLMHead(
+ config.vocab_size,
+ config.hidden_size,
+ prefix=maybe_prefix(prefix, "lm_head"),
+ )
+ else:
+ self.lm_head = PPMissingLayer()
+
+ self.logits_processor = LogitsProcessor(config.vocab_size)
+
+ def embed_input_ids(
+ self,
+ input_ids: torch.Tensor,
+ multimodal_embeddings: MultiModalEmbeddings | None = None,
+ *,
+ is_multimodal: torch.Tensor | None = None,
+ handle_oov_mm_token: bool = False,
+ ) -> torch.Tensor:
+ inputs_embeds = self._embed_text_input_ids(
+ input_ids,
+ self.model.embed_input_ids,
+ is_multimodal=is_multimodal,
+ handle_oov_mm_token=handle_oov_mm_token,
+ )
+
+ if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+ return inputs_embeds
+
+ is_multimodal = _require_is_multimodal(is_multimodal)
+
+ inputs_embeds = _merge_multimodal_embeddings(
+ inputs_embeds=inputs_embeds,
+ multimodal_embeddings=multimodal_embeddings,
+ is_multimodal=is_multimodal,
+ )
+
+ return inputs_embeds
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ inputs_embeds: torch.Tensor | None = None,
+ input_ids: torch.Tensor | None = None,
+ positions: torch.Tensor | None = None,
+ intermediate_tensors: IntermediateTensors | None = None,
+ **kwargs: object,
+ ):
+ draft_tokens = []
+ logits_lst = []
+ current_inputs_embeds = inputs_embeds
+ if current_inputs_embeds is None:
+ # [DEBUG] 打印词表大小和 input_ids 的范围
+ '''vocab_size = self.model.vocab_size
+ max_id = input_ids.max().item() if input_ids is not None else None
+ min_id = input_ids.min().item() if input_ids is not None else None
+ print(f"[MEDUSA DEBUG] embed_tokens vocab_size={vocab_size}, "
+ f"input_ids shape={input_ids.shape if input_ids is not None else None}, "
+ f"input_ids min={min_id}, max={max_id}")
+ if max_id is not None and max_id >= vocab_size:
+ print(f"[MEDUSA DEBUG] WARNING: input_ids max ({max_id}) >= vocab_size ({vocab_size})! "
+ f"This will cause index out of bounds!")'''
+ current_inputs_embeds = self.model.embed_tokens(input_ids)
+ current_hidden_states = hidden_states
+ for i in range(self.num_mtp_tokens):
+ next_hidden = self.model(
+ input_ids=input_ids,
+ positions=positions,
+ hidden_states=current_hidden_states,
+ intermediate_tensors=intermediate_tensors,
+ inputs_embeds=current_inputs_embeds
+ )
+ next_logits = self.logits_processor(self.lm_head, next_hidden)
+ # logits_lst.append(next_logits) # 保存
+
+ # 4. 采样下一个 Token (为了进行下一次递归 t+2 -> t+3)
+ # Only sample from the last position (the actual sampling position)
+ current_token_ids = torch.argmax(next_logits, dim=-1)
+ draft_tokens.append(current_token_ids)
+
+ # 5. 更新 Hidden State 进入下一次迭代
+ # Only use the last position for next iteration
+ current_hidden_states = next_hidden
+ current_inputs_embeds = self.model.embed_tokens(current_token_ids)
+ draft_tokens = torch.stack(draft_tokens, dim=1)
+ return draft_tokens
+ # return logits_lst
+
+ # def compute_logits(
+ # self,
+ # hidden_states: torch.Tensor,
+ # spec_step_idx: int = 0,
+ # ) -> torch.Tensor | None:
+ # return self.logits_processor(self.lm_head, hidden_states)
+
+ def compute_logits(
+ self,
+ hidden_states: list[torch.Tensor],
+ ) -> list[torch.Tensor]:
+ """
+ Compute logits using the SHARED lm_head.
+ """
+ logits_lst: list[torch.Tensor] = []
+ for hs in hidden_states:
+ _logits = self.logits_processor(self.lm_head, hs)
+ logits_lst.append(_logits)
+ return logits_lst
+
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ def remap_weight_names(weights):
+ for name, weight in weights:
+ # Strip language_model. prefix (InternVL checkpoint format)
+ if name.startswith("language_model."):
+ name = name[len("language_model."):]
+ if name.startswith("mtp."):
+ name = name.replace("mtp.", "model.")
+ elif any(key in name for key in ["embed_tokens", "lm_head"]):
+ pass # already correct after stripping language_model.
+ else:
+ continue
+ yield name, weight
+
+ loader = AutoWeightsLoader(self)
+ return loader.load_weights(remap_weight_names(weights))
diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py
index ee2bb837a498..f53a0e9bc629 100644
--- a/vllm/model_executor/models/qwen2_5_omni_thinker.py
+++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py
@@ -353,6 +353,39 @@ def get_target_channels(self) -> int:
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
return {"audio": None, "image": None, "video": None}
+ def get_mm_max_tokens_per_item(
+ self,
+ seq_len: int,
+ mm_counts: Mapping[str, int] | None = None,
+ ) -> Mapping[str, int] | None:
+ mm_counts = mm_counts or {}
+ requested_modalities = {m for m, c in mm_counts.items() if c > 0}
+ mm_max_tokens: dict[str, int] = {}
+
+ if requested_modalities & {"image", "video"}:
+ vl_tokens = Qwen2_5_VLProcessingInfo.get_mm_max_tokens_per_item(
+ self,
+ seq_len=seq_len,
+ mm_counts=mm_counts,
+ )
+ mm_max_tokens.update(
+ {
+ m: vl_tokens[m]
+ for m in ["image", "video"]
+ if m in requested_modalities
+ }
+ )
+
+ if "audio" in requested_modalities:
+ audio_tokens = Qwen2AudioProcessingInfo.get_mm_max_tokens_per_item(
+ self,
+ seq_len=seq_len,
+ mm_counts=mm_counts,
+ )
+ mm_max_tokens["audio"] = audio_tokens["audio"]
+
+ return mm_max_tokens
+
class Qwen2_5OmniThinkerDummyInputsBuilder(
BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]
diff --git a/vllm/model_executor/models/qwen2_audio.py b/vllm/model_executor/models/qwen2_audio.py
index 053e8bb85a3f..d125570a1122 100644
--- a/vllm/model_executor/models/qwen2_audio.py
+++ b/vllm/model_executor/models/qwen2_audio.py
@@ -179,6 +179,26 @@ def get_target_channels(self) -> int:
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
return {"audio": None}
+ def get_mm_max_tokens_per_item(
+ self,
+ seq_len: int,
+ mm_counts: Mapping[str, int] | None = None,
+ ) -> Mapping[str, int]:
+ mm_counts = mm_counts or {}
+ if mm_counts.get("audio", 0) <= 0:
+ return {}
+
+ feature_extractor = self.get_feature_extractor()
+ chunk_length = min(feature_extractor.chunk_length, 30)
+ audio_len = int(chunk_length * feature_extractor.sampling_rate)
+ hop_length = feature_extractor.hop_length
+ max_mel_seq_len = audio_len // hop_length
+
+ input_lengths = torch.tensor([max_mel_seq_len], dtype=torch.long)
+ _, output_lengths = _get_feat_extract_output_lengths(input_lengths)
+
+ return {"audio": int(output_lengths.item())}
+
class Qwen2AudioDummyInputsBuilder(BaseDummyInputsBuilder[Qwen2AudioProcessingInfo]):
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py
index 266ad5477b33..575ef8715a60 100644
--- a/vllm/model_executor/models/qwen3.py
+++ b/vllm/model_executor/models/qwen3.py
@@ -339,6 +339,6 @@ def compute_logits(
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
- skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
+ skip_prefixes=(["lm_head.", 'mtp'] if self.config.tie_word_embeddings else ['mtp']),
)
return loader.load_weights(weights)
diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py
index 1e6348b7218d..00335b88b137 100755
--- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py
+++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py
@@ -1163,6 +1163,39 @@ def get_feature_extractor(self, **kwargs: object):
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
return {"audio": None, "image": None, "video": None}
+ def get_mm_max_tokens_per_item(
+ self,
+ seq_len: int,
+ mm_counts: Mapping[str, int] | None = None,
+ ) -> Mapping[str, int] | None:
+ mm_counts = mm_counts or {}
+ requested_modalities = {m for m, c in mm_counts.items() if c > 0}
+ mm_max_tokens: dict[str, int] = {}
+
+ if requested_modalities & {"image", "video"}:
+ vl_tokens = Qwen2_5_VLProcessingInfo.get_mm_max_tokens_per_item(
+ self,
+ seq_len=seq_len,
+ mm_counts=mm_counts,
+ )
+ mm_max_tokens.update(
+ {
+ m: vl_tokens[m]
+ for m in ["image", "video"]
+ if m in requested_modalities
+ }
+ )
+
+ if "audio" in requested_modalities:
+ audio_tokens = Qwen2AudioProcessingInfo.get_mm_max_tokens_per_item(
+ self,
+ seq_len=seq_len,
+ mm_counts=mm_counts,
+ )
+ mm_max_tokens["audio"] = audio_tokens["audio"]
+
+ return mm_max_tokens
+
Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder
diff --git a/vllm/model_executor/models/qwen_vl.py b/vllm/model_executor/models/qwen_vl.py
index 8ac541f73d8b..1eb8ecc2d2aa 100644
--- a/vllm/model_executor/models/qwen_vl.py
+++ b/vllm/model_executor/models/qwen_vl.py
@@ -6,11 +6,9 @@
# Copyright (c) Alibaba Cloud.
"""Inference-only Qwen-VL model compatible with HuggingFace weights."""
-import copy
import math
-import unicodedata
-from collections.abc import Callable, Collection, Mapping, Sequence, Set
-from functools import lru_cache, partial
+from collections.abc import Callable, Mapping, Sequence
+from functools import partial
from typing import Annotated, Literal, TypeAlias
import regex as re
@@ -436,60 +434,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
)
-@lru_cache(maxsize=1)
-def _get_tokenizer_without_image_pad(
- tokenizer: PreTrainedTokenizer,
-) -> PreTrainedTokenizer:
- """
- The logic of adding image pad tokens should only be applied in
- [`QwenVLProcessor`][vllm.model_executor.models.qwen_vl.QwenVLProcessor],
- so they are patched out here.
-
- The definition of the wrapped tokenizer can be found here:
- https://huggingface.co/Qwen/Qwen-VL/blob/main/tokenization_qwen.py
- """
- new_tokenizer = copy.deepcopy(tokenizer)
-
- class TokenizerWithoutImagePad(tokenizer.__class__): # type: ignore
- def tokenize(
- self,
- text: str,
- allowed_special: Set[str] | str = "all",
- disallowed_special: Collection[str] | str = (),
- **kwargs,
- ) -> list[bytes | str]:
- text = unicodedata.normalize("NFC", text)
-
- return [
- self.decoder[t]
- for t in self.tokenizer.encode(
- text,
- allowed_special=allowed_special,
- disallowed_special=disallowed_special,
- )
- ]
-
- def _decode(
- self,
- token_ids: int | list[int],
- skip_special_tokens: bool = False,
- errors: str | None = None,
- **kwargs,
- ) -> str:
- if isinstance(token_ids, int):
- token_ids = [token_ids]
-
- return self.tokenizer.decode(
- token_ids,
- errors=errors or self.errors,
- )
-
- TokenizerWithoutImagePad.__name__ = f"{tokenizer.__class__.__name__}WithoutImagePad"
-
- new_tokenizer.__class__ = TokenizerWithoutImagePad
- return new_tokenizer
-
-
class QwenVLProcessor:
"""
This model doesn't define its own HF processor,
@@ -574,12 +518,6 @@ def __call__(
class QwenVLProcessingInfo(BaseProcessingInfo):
- def get_tokenizer(self) -> PreTrainedTokenizer:
- tokenizer = self.ctx.get_tokenizer()
- assert isinstance(tokenizer, PreTrainedTokenizer)
-
- return _get_tokenizer_without_image_pad(tokenizer)
-
def get_hf_processor(self, **kwargs: object) -> QwenVLProcessor:
return self.ctx.init_processor(
QwenVLProcessor,
diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py
index df75e8584f50..8c78db6f1878 100644
--- a/vllm/reasoning/__init__.py
+++ b/vllm/reasoning/__init__.py
@@ -68,6 +68,10 @@
"mistral_reasoning_parser",
"MistralReasoningParser",
),
+ "nemotron_v3": (
+ "nemotron_v3_reasoning_parser",
+ "NemotronV3ReasoningParser",
+ ),
"olmo3": (
"olmo3_reasoning_parser",
"Olmo3ReasoningParser",
diff --git a/vllm/reasoning/nemotron_v3_reasoning_parser.py b/vllm/reasoning/nemotron_v3_reasoning_parser.py
new file mode 100644
index 000000000000..2d3dc3685e98
--- /dev/null
+++ b/vllm/reasoning/nemotron_v3_reasoning_parser.py
@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm.entrypoints.openai.chat_completion.protocol import (
+ ChatCompletionRequest,
+)
+from vllm.entrypoints.openai.responses.protocol import (
+ ResponsesRequest,
+)
+from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser
+
+
+class NemotronV3ReasoningParser(DeepSeekR1ReasoningParser):
+ """
+ Reasoning parser for Nemotron V3 models.
+ """
+
+ def extract_reasoning(
+ self, model_output: str, request: ChatCompletionRequest | ResponsesRequest
+ ) -> tuple[str | None, str | None]:
+ reasoning_content, final_content = super().extract_reasoning(
+ model_output, request
+ )
+ chat_template_kwargs = getattr(request, "chat_template_kwargs", None)
+
+ if (
+ chat_template_kwargs
+ and (
+ chat_template_kwargs.get("enable_thinking") is False
+ or chat_template_kwargs.get("force_nonempty_content") is True
+ )
+ and final_content is None
+ ):
+ reasoning_content, final_content = final_content, reasoning_content
+
+ return reasoning_content, final_content
diff --git a/vllm/renderers/qwen_vl.py b/vllm/renderers/qwen_vl.py
new file mode 100644
index 000000000000..4b47d0216bfa
--- /dev/null
+++ b/vllm/renderers/qwen_vl.py
@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Any
+
+from vllm.config import VllmConfig
+from vllm.tokenizers import cached_get_tokenizer
+from vllm.tokenizers.qwen_vl import QwenVLTokenizer
+
+from .base import BaseRenderer
+from .hf import HfRenderer
+
+
+class QwenVLRenderer(BaseRenderer[QwenVLTokenizer]):
+ @classmethod
+ def from_config( # type: ignore[override]
+ cls,
+ config: VllmConfig,
+ tokenizer_kwargs: dict[str, Any],
+ ) -> "HfRenderer":
+ model_config = config.model_config
+ if model_config.skip_tokenizer_init:
+ tokenizer = None
+ else:
+ tokenizer = cached_get_tokenizer(
+ tokenizer_cls=QwenVLTokenizer,
+ **tokenizer_kwargs,
+ )
+
+ return HfRenderer(config, tokenizer)
diff --git a/vllm/renderers/registry.py b/vllm/renderers/registry.py
index cd09c80f9f95..de95505eca68 100644
--- a/vllm/renderers/registry.py
+++ b/vllm/renderers/registry.py
@@ -20,6 +20,7 @@
"hf": ("hf", "HfRenderer"),
"grok2": ("grok2", "Grok2Renderer"),
"mistral": ("mistral", "MistralRenderer"),
+ "qwen_vl": ("qwen_vl", "QwenVLRenderer"),
"terratorch": ("terratorch", "TerratorchRenderer"),
}
diff --git a/vllm/tokenizers/deepseek_v32.py b/vllm/tokenizers/deepseek_v32.py
index 28071ef6970c..4525eaa343c9 100644
--- a/vllm/tokenizers/deepseek_v32.py
+++ b/vllm/tokenizers/deepseek_v32.py
@@ -7,9 +7,9 @@
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
-from . import TokenizerLike
from .deepseek_v32_encoding import encode_messages
from .hf import HfTokenizer, get_cached_tokenizer
+from .protocol import TokenizerLike
def get_deepseek_v32_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer:
diff --git a/vllm/tokenizers/qwen_vl.py b/vllm/tokenizers/qwen_vl.py
new file mode 100644
index 000000000000..5b506df4df62
--- /dev/null
+++ b/vllm/tokenizers/qwen_vl.py
@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
+import unicodedata
+from collections.abc import Collection, Set
+
+from transformers import AutoTokenizer
+
+from .hf import HfTokenizer, get_cached_tokenizer
+from .protocol import TokenizerLike
+
+
+def get_qwen_vl_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer:
+ """
+ The logic of adding image pad tokens should only be applied in
+ `QwenVLProcessor`, so they are patched out here.
+
+ The definition of the wrapped tokenizer can be found here:
+ https://huggingface.co/Qwen/Qwen-VL/blob/main/tokenization_qwen.py
+ """
+ new_tokenizer = copy.copy(tokenizer)
+
+ class TokenizerWithoutImagePad(tokenizer.__class__): # type: ignore
+ def tokenize(
+ self,
+ text: str,
+ allowed_special: Set[str] | str = "all",
+ disallowed_special: Collection[str] | str = (),
+ **kwargs,
+ ) -> list[bytes | str]:
+ text = unicodedata.normalize("NFC", text)
+
+ return [
+ self.decoder[t]
+ for t in self.tokenizer.encode(
+ text,
+ allowed_special=allowed_special,
+ disallowed_special=disallowed_special,
+ )
+ ]
+
+ def _decode(
+ self,
+ token_ids: int | list[int],
+ skip_special_tokens: bool = False,
+ errors: str | None = None,
+ **kwargs,
+ ) -> str:
+ if isinstance(token_ids, int):
+ token_ids = [token_ids]
+
+ return self.tokenizer.decode(
+ token_ids,
+ errors=errors or self.errors,
+ )
+
+ TokenizerWithoutImagePad.__name__ = f"{tokenizer.__class__.__name__}WithoutImagePad"
+
+ new_tokenizer.__class__ = TokenizerWithoutImagePad
+ return new_tokenizer
+
+
+class QwenVLTokenizer(TokenizerLike):
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs) -> HfTokenizer:
+ tokenizer = AutoTokenizer.from_pretrained(*args, **kwargs)
+ return get_cached_tokenizer(get_qwen_vl_tokenizer(tokenizer))
diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py
index 2da7842b03d5..4512f766c99b 100644
--- a/vllm/tokenizers/registry.py
+++ b/vllm/tokenizers/registry.py
@@ -36,6 +36,7 @@
"grok2": ("grok2", "Grok2Tokenizer"),
"hf": ("hf", "CachedHfTokenizer"),
"mistral": ("mistral", "MistralTokenizer"),
+ "qwen_vl": ("qwen_vl", "QwenVLTokenizer"),
}
@@ -165,6 +166,10 @@ def resolve_tokenizer_args(
):
tokenizer_mode = "grok2"
+ # Model-specific tokenizers
+ if tokenizer_mode == "auto" and "/Qwen-VL" in str(tokenizer_name):
+ tokenizer_mode = "qwen_vl"
+
# Fallback to HF tokenizer
if tokenizer_mode == "auto":
tokenizer_mode = "hf"
diff --git a/vllm/utils/math_utils.py b/vllm/utils/math_utils.py
index a0e301af471f..1ea4401e1568 100644
--- a/vllm/utils/math_utils.py
+++ b/vllm/utils/math_utils.py
@@ -30,3 +30,8 @@ def round_up(x: int, y: int) -> int:
def round_down(x: int, y: int) -> int:
"""Round down x to the nearest multiple of y."""
return (x // y) * y
+
+
+def largest_power_of_2_divisor(n: int) -> int:
+ """Return the largest power-of-2 that divides *n* (isolate lowest set bit)."""
+ return n & (-n)
diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py
index 585ad1d793ff..e8d65bb99e74 100644
--- a/vllm/v1/attention/backend.py
+++ b/vllm/v1/attention/backend.py
@@ -86,6 +86,26 @@ def get_kv_cache_shape(
) -> tuple[int, ...]:
raise NotImplementedError
+ @classmethod
+ def get_kv_cache_block_dim(
+ cls,
+ block_size: int,
+ num_kv_heads: int,
+ head_size: int,
+ cache_dtype_str: str = "auto",
+ ) -> int:
+ """Discover which tensor dim is the block index, since different
+ backends lay out dims differently."""
+ _S = 1234567
+ shape = cls.get_kv_cache_shape(
+ _S,
+ block_size,
+ num_kv_heads,
+ head_size,
+ cache_dtype_str=cache_dtype_str,
+ )
+ return shape.index(_S)
+
@staticmethod
def get_kv_cache_stride_order(
include_num_layers_dimension: bool = False,
diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py
index e84312970989..a0ab157f5384 100644
--- a/vllm/v1/attention/backends/mla/indexer.py
+++ b/vllm/v1/attention/backends/mla/indexer.py
@@ -372,12 +372,14 @@ def build(
# [7, 6, 8, 0] -> [7, 7, 7, 6, 8, 8, 8, 8]
expanded_base = torch.repeat_interleave(
- seq_lens - decode_lens, decode_lens
+ seq_lens - decode_lens, decode_lens, output_size=actual_expanded
)
# [0, 3, 4, 8] -> [0, 0, 0, 3, 4, 4, 4, 4]
expanded_starts = torch.repeat_interleave(
- common_attn_metadata.query_start_loc[:num_decodes], decode_lens
+ common_attn_metadata.query_start_loc[:num_decodes],
+ decode_lens,
+ output_size=actual_expanded,
)
# [0, 1, 2, 0, 0, 1, 2, 3]
@@ -395,7 +397,9 @@ def build(
# Give each of the flattened entries the same block table row as the
# original request.
self.expanded_block_table_buffer[:actual_expanded] = (
- torch.repeat_interleave(block_table, decode_lens, dim=0)
+ torch.repeat_interleave(
+ block_table, decode_lens, dim=0, output_size=actual_expanded
+ )
)
if actual_expanded < num_decode_tokens:
self.expanded_block_table_buffer[
diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py
index eaa95dfe49f7..2fa1f978bda8 100644
--- a/vllm/v1/core/kv_cache_coordinator.py
+++ b/vllm/v1/core/kv_cache_coordinator.py
@@ -14,6 +14,7 @@
)
from vllm.v1.core.single_type_kv_cache_manager import (
CrossAttentionManager,
+ MambaManager,
SingleTypeKVCacheManager,
get_manager_for_kv_cache_spec,
)
@@ -198,6 +199,19 @@ def free(self, request_id: str) -> None:
for manager in self.single_type_managers:
manager.free(request_id)
+ def free_tail_blocks(self, request_id: str,
+ num_tokens_to_keep: int) -> None:
+ """Free tail blocks for a request across all KV cache groups,
+ keeping blocks that cover at least num_tokens_to_keep tokens.
+
+ Args:
+ request_id: The request ID.
+ num_tokens_to_keep: Number of prefix tokens to retain.
+ """
+ for manager in self.single_type_managers:
+ num_blocks_to_keep = num_tokens_to_keep // manager.block_size
+ manager.free_tail_blocks(request_id, num_blocks_to_keep)
+
def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
"""
Get the number of common prefix blocks for all requests with allocated
@@ -239,6 +253,73 @@ def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]:
for manager in self.single_type_managers
)
+ def fork_blocks(
+ self,
+ parent_request_id: str,
+ child_request_id: str,
+ num_blocks_to_share: int,
+ ) -> tuple[list[KVCacheBlock], ...]:
+ """
+ Fork blocks from parent request to child request.
+
+ For Attention layers: shares the parent's blocks with the child by
+ increasing reference counts, avoiding actual data copy.
+
+ For Mamba layers: does NOT share blocks because Mamba state is cumulative
+ and would cause pollution between parent and child branches. The child
+ will allocate new blocks and the Mamba state will be copied separately
+ at runtime via mamba_state_copy_func.
+
+ Args:
+ parent_request_id: The parent request ID.
+ child_request_id: The child request ID.
+ num_blocks_to_share: Number of blocks to share from parent.
+
+ Returns:
+ The shared blocks for the child request (empty for Mamba managers).
+ """
+ shared_blocks_per_group: list[list[KVCacheBlock]] = []
+
+ for manager in self.single_type_managers:
+ # IMPORTANT: Skip sharing for MambaManager
+ # Mamba state is cumulative/stateful, so parent and child would
+ # pollute each other's state if they shared the same blocks.
+ # The child will get new blocks allocated during scheduling,
+ # and the Mamba state copy will happen at GPU model runner level.
+ if isinstance(manager, MambaManager):
+ shared_blocks_per_group.append([])
+ # Initialize empty block list for child so it gets fresh blocks
+ # NOTE: Do NOT set num_cached_block here! The child request is not
+ # "running" yet - it will go through the normal scheduling path
+ # which may find prefix cache hits (new_computed_blocks).
+ # Setting num_cached_block would make get_num_blocks_to_allocate()
+ # think this is a running request and assert new_computed_blocks == 0.
+ manager.req_to_blocks[child_request_id] = []
+ continue
+
+ parent_blocks = manager.req_to_blocks.get(parent_request_id)
+ if parent_blocks is None:
+ shared_blocks_per_group.append([])
+ continue
+
+ # Share blocks up to num_blocks_to_share
+ blocks_to_share = parent_blocks[:num_blocks_to_share]
+ shared_blocks: list[KVCacheBlock] = []
+
+ for block in blocks_to_share:
+ # Increase reference count for shared block
+ block.ref_cnt += 1
+ shared_blocks.append(block)
+
+ # Register the shared blocks for the child request
+ manager.req_to_blocks[child_request_id] = shared_blocks.copy()
+ # IMPORTANT: Also set num_cached_block to avoid assertion error
+ # in save_new_computed_blocks when the child is scheduled
+ manager.num_cached_block[child_request_id] = len(shared_blocks)
+ shared_blocks_per_group.append(shared_blocks)
+
+ return tuple(shared_blocks_per_group)
+
@abstractmethod
def find_longest_cache_hit(
self,
diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py
index 7f8d804753d2..6bb2875fb426 100644
--- a/vllm/v1/core/kv_cache_manager.py
+++ b/vllm/v1/core/kv_cache_manager.py
@@ -385,6 +385,37 @@ def free(self, request: Request) -> None:
"""
self.coordinator.free(request.request_id)
+ def free_partial(self, request: Request,
+ num_tokens_to_keep: int) -> int:
+ """Free tail blocks for a request, keeping blocks that cover
+ num_tokens_to_keep tokens. The kept token count is aligned down
+ to the block boundary.
+
+ Args:
+ request: The request to partially free.
+ num_tokens_to_keep: Desired number of prefix tokens to retain.
+
+ Returns:
+ The actual number of tokens retained (block-aligned).
+ """
+ if num_tokens_to_keep <= 0:
+ self.coordinator.free(request.request_id)
+ return 0
+
+ # Use the first manager's block_size for alignment
+ # (all managers cover the same token positions)
+ first_manager = self.coordinator.single_type_managers[0]
+ block_size = first_manager.block_size
+ num_blocks_to_keep = num_tokens_to_keep // block_size
+
+ if num_blocks_to_keep <= 0:
+ self.coordinator.free(request.request_id)
+ return 0
+
+ self.coordinator.free_tail_blocks(
+ request.request_id, num_tokens_to_keep)
+ return num_blocks_to_keep * block_size
+
def remove_skipped_blocks(
self, request_id: str, total_computed_tokens: int
) -> None:
@@ -489,6 +520,79 @@ def create_kv_cache_blocks(
# Only create new KVCacheBlocks for non-empty blocks
return KVCacheBlocks(blocks) if any(blocks) else self.empty_kv_cache_blocks
+ def take_new_block_ids(self) -> list[int]:
+ """Drain and return new attention block IDs for zeroing."""
+ ids: list[int] = []
+ for mgr in self.coordinator.single_type_managers:
+ ids.extend(mgr.take_new_block_ids())
+ return ids
+
def new_step_starts(self) -> None:
"""Called when a new step is started."""
self.coordinator.new_step_starts()
+
+ def fork_kv_cache(
+ self,
+ parent_request_id: str,
+ child_request_id: str,
+ fork_position: int,
+ block_size: int, # Add block_size parameter
+ ) -> KVCacheBlocks:
+ """Fork KV cache from parent request to child request.
+
+ This method shares ONLY the fully filled blocks from parent to child.
+ The last partial block is NOT shared to avoid KV cache corruption when
+ parent and child write to the same block position with different KV values.
+
+ Args:
+ parent_request_id: The parent request ID to fork from.
+ child_request_id: The child request ID to fork to.
+ fork_position: The token position where fork occurred (exclusive).
+ This is the position of the FORK/CHILD token.
+ block_size: The block size to use for calculating number of blocks.
+
+ Returns:
+ KVCacheBlocks: The shared blocks for the child request.
+ """
+ debug = False
+ # IMPORTANT: Only share FULLY FILLED blocks
+ # If fork_position = 2117 and block_size = 16:
+ # - num_full_blocks = 2117 // 16 = 132 (blocks 0-131 are fully filled)
+ # - Block 132 has positions 2112-2117 (6 tokens), NOT fully filled
+ # - We only share blocks 0-131, NOT block 132
+ #
+ # This prevents the bug where:
+ # - Parent writes FORK token KV at block 132, position 5
+ # - Child writes CHILD token KV at block 132, position 5
+ # - They overwrite each other's KV!
+ num_full_blocks = fork_position // block_size
+ if debug:
+ logger.info(
+ "Fork KV cache: parent=%s, child=%s, fork_position=%d, "
+ "block_size=%d, num_full_blocks_to_share=%d",
+ parent_request_id, child_request_id, fork_position,
+ block_size, num_full_blocks
+ )
+
+ if num_full_blocks == 0:
+ # No full blocks to share, child will compute everything from scratch
+ if debug:
+ logger.info("No full blocks to share, child will recompute all KV")
+ return self.empty_kv_cache_blocks
+
+ # Use coordinator to share ONLY the full blocks between parent and child
+ shared_blocks = self.coordinator.fork_blocks(
+ parent_request_id=parent_request_id,
+ child_request_id=child_request_id,
+ num_blocks_to_share=num_full_blocks, # Only full blocks!
+ )
+ if debug:
+ logger.info(
+ "Forked KV cache from %s to %s at position %d, "
+ "shared %d full blocks (tokens 0-%d are shared, tokens %d+ will be recomputed)",
+ parent_request_id, child_request_id, fork_position,
+ num_full_blocks, num_full_blocks * block_size - 1,
+ num_full_blocks * block_size
+ )
+
+ return self.create_kv_cache_blocks(shared_blocks)
diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py
index 0b3958dbcf5a..962d67e0a698 100644
--- a/vllm/v1/core/sched/async_scheduler.py
+++ b/vllm/v1/core/sched/async_scheduler.py
@@ -36,15 +36,15 @@ def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
def _update_request_with_output(
self, request: Request, new_token_ids: list[int]
- ) -> tuple[list[int], bool]:
+ ) -> tuple[list[int], bool, bool]:
if request.discard_latest_async_tokens:
# If the request is force preempted in reset_prefix_cache, we
# should discard the latest async token.
request.discard_latest_async_tokens = False
- return [], False
+ return [], False, False
status_before_update = request.status
- new_token_ids, stopped = super()._update_request_with_output(
+ new_token_ids, stopped, forked = super()._update_request_with_output(
request, new_token_ids
)
@@ -57,4 +57,4 @@ def _update_request_with_output(
self.kv_cache_manager.cache_blocks(
request, request.num_computed_tokens - request.num_output_placeholders
)
- return new_token_ids, stopped
+ return new_token_ids, stopped, forked
diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py
index 0f6ac98fdc15..556a01b3b500 100644
--- a/vllm/v1/core/sched/output.py
+++ b/vllm/v1/core/sched/output.py
@@ -233,6 +233,17 @@ class SchedulerOutput:
# EC Cache Connector metadata
ec_connector_metadata: ECConnectorMetadata | None = None
+ # Block IDs freshly allocated from the pool during this scheduling step.
+ # The worker zeros the corresponding GPU memory before the blocks are used,
+ # preventing stale NaN/data from corrupting attention or SSM computation.
+ new_block_ids_to_zero: list[int] | None = None
+
+ # Fork information: child_req_id -> parent_req_id
+ # Used to copy Mamba state from parent to forked child requests.
+ # When a child request is scheduled for the first time, its Mamba state
+ # needs to be initialized from the parent's Mamba state.
+ fork_child_to_parent: dict[str, str] | None = None
+
@classmethod
def make_empty(cls) -> "SchedulerOutput":
return cls(
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index bf397ad681ca..e4145ba250a2 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -45,13 +45,17 @@
NewRequestData,
SchedulerOutput,
)
-from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
+from vllm.v1.core.sched.request_queue import (
+ FCFSRequestQueue,
+ SchedulingPolicy,
+ create_request_queue,
+)
from vllm.v1.core.sched.utils import check_stop, remove_all
from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
-from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
+from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.metrics.perf import ModelMetrics, PerfStats
from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats
-from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
+from vllm.v1.outputs import DraftTokenIds, ForkInfo, KVConnectorOutput, ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
from vllm.v1.spec_decode.metrics import SpecDecodingStats
from vllm.v1.structured_output import StructuredOutputManager
@@ -88,6 +92,58 @@ def __init__(
self.structured_output_manager = structured_output_manager
self.is_encoder_decoder = vllm_config.model_config.is_encoder_decoder
+ # Fork-related configuration
+ # Read from HuggingFace config (config.json), not vLLM's ModelConfig
+ hf_config = getattr(vllm_config.model_config, 'hf_config', None)
+ if hf_config is not None:
+ self.fork_token_id: int | None = getattr(hf_config, 'fork_token_id', None)
+ self.child_token_id: int | None = getattr(hf_config, 'child_token_id', None)
+ else:
+ self.fork_token_id = None
+ self.child_token_id = None
+
+
+ self.enable_dynamic_fork = (self.fork_token_id is not None and
+ self.child_token_id is not None)
+ # self.enable_dynamic_fork = False
+ # Fork debug switch - controlled by environment variable VLLM_FORK_DEBUG
+ import os
+ self.fork_debug = os.environ.get("VLLM_FORK_DEBUG", "0").lower() in ("1", "true", "yes")
+
+ if self.enable_dynamic_fork:
+ logger.info(f"Dynamic fork enabled: fork_token_id={self.fork_token_id}, "
+ f"child_token_id={self.child_token_id}, debug={self.fork_debug}")
+ # Only load tokenizer when debug is enabled (performance optimization)
+ if self.fork_debug:
+ from transformers import AutoTokenizer
+ model_path = vllm_config.model_config.model
+ self.fork_tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+ logger.info(f"Loaded tokenizer from {model_path} for fork debugging")
+ else:
+ self.fork_tokenizer = None
+ else:
+ self.fork_tokenizer = None
+
+ # Track parent-child relationships for result aggregation
+ self.fork_relationships: dict[str, list[str]] = {} # parent_id -> [child_ids]
+ self.pending_fork_requests: list[Request] = [] # Requests to add after current step
+
+ # Fork output aggregation
+ # Stores accumulated token IDs for each fork branch
+ self.fork_outputs: dict[str, dict[str, list[int]]] = {} # parent_id -> {req_id: [token_ids]}
+ # Tracks which branches have finished
+ self.fork_finished: dict[str, set[str]] = {} # parent_id -> set of finished req_ids
+ # Stores the original parent request's client_index for returning aggregated results
+ self.fork_client_index: dict[str, int] = {} # parent_id -> client_index
+
+ # Track fork child->parent mapping for Mamba state copy
+ # This is populated in _process_pending_fork_requests and cleared after each schedule step
+ self.fork_child_to_parent_this_step: dict[str, str] = {} # child_id -> parent_id
+
+ # Dirty flag for incremental waiting queue sorting:
+ # Only re-sort when fork events actually change the queue composition.
+ self._fork_queue_dirty: bool = False
+
# include_finished_set controls whether a separate set of finished
# request ids should be included in the EngineCoreOutputs returned
# by update_from_outputs(). This is currently used in the multi-engine
@@ -104,6 +160,7 @@ def __init__(
if self.scheduler_config.max_num_scheduled_tokens
else self.scheduler_config.max_num_batched_tokens
)
+ print(self.max_num_scheduled_tokens)
self.max_model_len = vllm_config.model_config.max_model_len
self.enable_kv_cache_events = (
self.kv_events_config is not None
@@ -233,13 +290,8 @@ def __init__(
self.use_pp = self.parallel_config.pipeline_parallel_size > 1
self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
- def has_mamba_layers(kv_cache_config: KVCacheConfig) -> bool:
- return any(
- isinstance(group_spec.kv_cache_spec, MambaSpec)
- for group_spec in kv_cache_config.kv_cache_groups
- )
-
- self.has_mamba_layers = has_mamba_layers(kv_cache_config)
+ self.has_mamba_layers = kv_cache_config.has_mamba_layers
+ self.needs_kv_cache_zeroing = kv_cache_config.needs_kv_cache_zeroing
self.need_mamba_block_aligned_split = (
self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align"
)
@@ -269,6 +321,149 @@ def has_mamba_layers(kv_cache_config: KVCacheConfig) -> bool:
self._pause_state: PauseState = PauseState.UNPAUSED
+ # Fork-aware scheduling configuration
+ # KV cache usage threshold above which parent requests are not admitted
+ import os
+ self.fork_parent_kv_cache_threshold = float(
+ os.environ.get("VLLM_FORK_PARENT_KV_CACHE_THRESHOLD", "0.2"))
+ # Allow up to N fork-children to be admitted to running beyond
+ # max_num_running_reqs. This guarantees that a freshly-created child
+ # can be co-batched with its parent in the very next forward pass,
+ # preserving APAR's weight-load amortization benefit (Option A).
+ # Children share most of the parent's KV via fork_kv_cache so they
+ # only consume <=1 partial block + spec slots of new KV.
+ self.fork_child_running_overflow = int(
+ os.environ.get("VLLM_FORK_CHILD_RUNNING_OVERFLOW", "0"))
+ if self.enable_dynamic_fork:
+ logger.info(
+ "Fork-aware scheduling enabled: "
+ "parent KV cache admission threshold=%.2f, "
+ "child running overflow slots=%d",
+ self.fork_parent_kv_cache_threshold,
+ self.fork_child_running_overflow)
+
+ def _get_fork_waiting_priority(
+ self,
+ request: Request,
+ family_waiting_count: dict[str, int],
+ family_running_count: dict[str, int],
+ ) -> tuple:
+ """Compute a sort key for fork-aware waiting queue ordering.
+
+ Priority rules (FCFS-by-family):
+ - Child requests always have higher priority than parent requests.
+ - Among children: families whose parent arrived earlier get higher
+ priority (FCFS by family). Within the same family, more siblings
+ already running/finished = higher priority.
+ - Among parents/non-fork: earlier arrival = higher priority (FCFS).
+
+ Returns a tuple for sorting (lower = higher priority):
+ (type_order, ..., arrival_time, request_id)
+ """
+ if request.is_fork_child:
+ # Child request: type_order=0 (higher priority than parents)
+ # Primary sub-key: parent's arrival time (earlier parent = higher priority)
+ # Secondary sub-key: prefer families closer to completion
+ parent_id = request.parent_request_id
+ total_children = len(self.fork_relationships.get(parent_id, []))
+ waiting = family_waiting_count.get(parent_id, 0)
+ progress = total_children - waiting # running + finished siblings
+
+ parent_req = self.requests.get(parent_id)
+ parent_arrival = parent_req.arrival_time if parent_req else request.arrival_time
+
+ return (0, parent_arrival, -progress, request.arrival_time, request.request_id)
+ else:
+ # Parent request (or non-fork request): type_order=1 (lower priority)
+ # Pure FCFS: earlier arrival = higher priority
+ return (1, request.arrival_time, request.request_id)
+
+ def _sort_waiting_queue_for_fork(self) -> None:
+ """Re-sort the waiting queue using fork-aware priority rules.
+
+ Called at the beginning of schedule() when dynamic fork is enabled.
+ Uses a dirty flag to skip re-sorting when no fork events have occurred
+ since the last sort, reducing CPU overhead from O(n log n) to O(1) for
+ the common case where the queue composition hasn't changed.
+
+ The dirty flag is set by:
+ - _process_pending_fork_requests: new child requests added to waiting
+ - _collect_fork_output: a branch finished, changing family statistics
+ """
+ if not self.waiting or not self._fork_queue_dirty:
+ return
+
+ # Pre-compute family counts in O(n) instead of O(n²)
+ family_waiting_count: dict[str, int] = defaultdict(int)
+ for r in self.waiting:
+ if r.is_fork_child:
+ family_waiting_count[r.parent_request_id] += 1
+
+ family_running_count: dict[str, int] = defaultdict(int)
+ for r in self.running:
+ if r.is_fork_child:
+ family_running_count[r.parent_request_id] += 1
+
+ # Extract all requests, sort by fork-aware priority, rebuild queue
+ all_requests = list(self.waiting)
+ # Clear the existing queue
+ if isinstance(self.waiting, FCFSRequestQueue):
+ self.waiting.clear()
+ else:
+ # For PriorityRequestQueue
+ self.waiting._heap.clear()
+
+ all_requests.sort(
+ key=lambda r: self._get_fork_waiting_priority(
+ r, family_waiting_count, family_running_count))
+ for req in all_requests:
+ if isinstance(self.waiting, FCFSRequestQueue):
+ self.waiting.append(req)
+ else:
+ self.waiting.add_request(req)
+
+ self._fork_queue_dirty = False
+
+ def _select_preempt_request_for_fork(self) -> Request:
+ """Select a request to preempt from the running queue using fork-aware rules.
+
+ Preemption rules (FCFS-by-family, consistent with scheduling priority):
+ - Always prefer to preempt parent requests over child requests.
+ - Among parents: later arrival = preempted first (FCFS protection).
+ - Among non-fork requests: later arrival = preempted first.
+ - Among children (only if no parents/others): later parent arrival =
+ preempted first; ties broken by default priority/arrival_time.
+
+ Returns the request to preempt.
+ """
+ # Separate parents and children in running queue
+ parent_requests = []
+ child_requests = []
+ other_requests = []
+
+ for req in self.running:
+ if req.is_fork_child:
+ child_requests.append(req)
+ elif req.fork_count > 0:
+ # This is a parent that has produced children
+ parent_requests.append(req)
+ else:
+ other_requests.append(req)
+
+ if parent_requests:
+ # Preempt the parent that arrived latest (FCFS: earlier = protected)
+ return max(parent_requests, key=lambda r: r.arrival_time)
+ elif other_requests:
+ # Non-fork requests: preempt latest arrival
+ return max(other_requests, key=lambda r: r.arrival_time)
+ else:
+ # Only children remain: preempt children from later-arriving families first
+ def _child_preempt_key(r):
+ parent_req = self.requests.get(r.parent_request_id)
+ parent_arrival = parent_req.arrival_time if parent_req else r.arrival_time
+ return (-parent_arrival, r.priority, r.arrival_time)
+ return max(child_requests, key=_child_preempt_key)
+
def _mamba_block_aligned_split(
self,
request: Request,
@@ -354,6 +549,10 @@ def schedule(self) -> SchedulerOutput:
self.kv_cache_manager.new_step_starts()
+ # Fork-aware: re-sort waiting queue before scheduling
+ if self.enable_dynamic_fork:
+ self._sort_waiting_queue_for_fork()
+
# First, schedule the RUNNING requests.
req_index = 0
while req_index < len(self.running) and token_budget > 0:
@@ -443,10 +642,43 @@ def schedule(self) -> SchedulerOutput:
if new_blocks is not None:
# The request can be scheduled.
break
-
+ logger.info("preempt triggered: req=%s, kv_usage=%.3f",
+ request.request_id, self.kv_cache_manager.usage)
# The request cannot be scheduled.
# Preempt the lowest-priority request.
- if self.policy == SchedulingPolicy.PRIORITY:
+ # P0: preserve already-scheduled draft (spec) tokens of the
+ # preempted request so the MTP investment is not wasted.
+ # We re-attach them to `spec_token_ids` AFTER
+ # `_preempt_request` (which clears that field), so the
+ # next schedule step can verify them again.
+ preempted_spec_token_ids: list[int] | None = None
+ if self.enable_dynamic_fork:
+ # Fork-aware preemption: prefer parents over children
+ preempted_req = self._select_preempt_request_for_fork()
+ self.running.remove(preempted_req)
+ if preempted_req in scheduled_running_reqs:
+ preempted_req_id = preempted_req.request_id
+ scheduled_running_reqs.remove(preempted_req)
+ token_budget += num_scheduled_tokens.pop(preempted_req_id)
+ req_to_new_blocks.pop(preempted_req_id)
+ preempted_spec_token_ids = (
+ scheduled_spec_decode_tokens.pop(
+ preempted_req_id, None
+ )
+ )
+ logger.info("popped spec_tokens for %s: %s",
+ preempted_req_id, preempted_spec_token_ids)
+ preempted_encoder_inputs = scheduled_encoder_inputs.pop(
+ preempted_req_id, None
+ )
+ if preempted_encoder_inputs:
+ num_embeds_to_restore = sum(
+ preempted_req.get_num_encoder_embeds(i)
+ for i in preempted_encoder_inputs
+ )
+ encoder_compute_budget += num_embeds_to_restore
+ req_index -= 1
+ elif self.policy == SchedulingPolicy.PRIORITY:
preempted_req = max(
self.running,
key=lambda r: (r.priority, r.arrival_time),
@@ -457,7 +689,11 @@ def schedule(self) -> SchedulerOutput:
scheduled_running_reqs.remove(preempted_req)
token_budget += num_scheduled_tokens.pop(preempted_req_id)
req_to_new_blocks.pop(preempted_req_id)
- scheduled_spec_decode_tokens.pop(preempted_req_id, None)
+ preempted_spec_token_ids = (
+ scheduled_spec_decode_tokens.pop(
+ preempted_req_id, None
+ )
+ )
preempted_encoder_inputs = scheduled_encoder_inputs.pop(
preempted_req_id, None
)
@@ -474,6 +710,15 @@ def schedule(self) -> SchedulerOutput:
preempted_req = self.running.pop()
self._preempt_request(preempted_req, scheduled_timestamp)
+ # Restore the preserved draft tokens AFTER _preempt_request
+ # (which clears `spec_token_ids` internally). They remain
+ # valid as next-token predictions because preemption keeps
+ # the request's output tokens; only KV is recomputed.
+ '''if preempted_spec_token_ids:
+ print("preempted_spec_token_ids: ", preempted_spec_token_ids)
+ preempted_req.spec_token_ids = list(
+ preempted_spec_token_ids
+ )'''
preempted_reqs.append(preempted_req)
if preempted_req == request:
# No more request to preempt. Cannot schedule this request.
@@ -533,18 +778,63 @@ def schedule(self) -> SchedulerOutput:
assert len(scheduled_loras) <= self.lora_config.max_loras
# Next, schedule the WAITING requests.
- if not preempted_reqs and self._pause_state == PauseState.UNPAUSED:
+ # Fork-aware (Option A): even if preemption happened during the running
+ # loop, still allow fork-children to be admitted so that parent and
+ # child stay co-batched in the same forward pass. Non-child requests
+ # remain gated by the original `not preempted_reqs` rule.
+ fork_only_admission = (
+ self.enable_dynamic_fork
+ and bool(preempted_reqs)
+ )
+ if (
+ (not preempted_reqs or fork_only_admission)
+ and self._pause_state == PauseState.UNPAUSED
+ ):
# Use a temporary RequestQueue to collect requests that need to be
# skipped and put back at the head of the waiting queue later
skipped_waiting_requests = create_request_queue(self.policy)
+ # Fork-aware effective running cap: children may overflow up to
+ # `fork_child_running_overflow` slots above max_num_running_reqs.
+ # Non-children must still respect max_num_running_reqs.
+ effective_max_running = self.max_num_running_reqs
+ if self.enable_dynamic_fork:
+ effective_max_running += self.fork_child_running_overflow
+
while self.waiting and token_budget > 0:
- if len(self.running) == self.max_num_running_reqs:
+ if len(self.running) >= effective_max_running:
break
request = self.waiting.peek_request()
request_id = request.request_id
+ # Fork-only admission mode: skip non-child requests (they will
+ # be retried next step when preemption pressure has cleared).
+ if fork_only_admission and not request.is_fork_child:
+ self.waiting.pop_request()
+ skipped_waiting_requests.prepend_request(request)
+ continue
+
+ # Respect the original max_num_running_reqs cap for non-children
+ # even when overflow is enabled. Only children may use the
+ # overflow slots.
+ if (
+ not request.is_fork_child
+ and len(self.running) >= self.max_num_running_reqs
+ ):
+ self.waiting.pop_request()
+ skipped_waiting_requests.prepend_request(request)
+ continue
+
+ # Fork-aware: block parent requests when KV cache usage is high
+ if (self.enable_dynamic_fork
+ and not request.is_fork_child
+ and self.kv_cache_manager.usage
+ > self.fork_parent_kv_cache_threshold):
+ self.waiting.pop_request()
+ skipped_waiting_requests.prepend_request(request)
+ continue
+
# KVTransfer: skip request if still waiting for remote kvs.
if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
is_ready = self._update_waiting_for_remote_kv(request)
@@ -637,8 +927,11 @@ def schedule(self) -> SchedulerOutput:
num_new_local_computed_tokens + num_external_computed_tokens
)
else:
- # KVTransfer: WAITING reqs have num_computed_tokens > 0
- # after async KV recvs are completed.
+ # num_computed_tokens > 0 for two cases:
+ # 1. KVTransfer: async KV recvs completed.
+ # 2. Partial preemption: prompt blocks still allocated.
+ # In both cases, blocks are already in req_to_blocks,
+ # so we only need to allocate the remaining tokens.
new_computed_blocks = self.kv_cache_manager.empty_kv_cache_blocks
num_new_local_computed_tokens = 0
num_computed_tokens = request.num_computed_tokens
@@ -656,7 +949,24 @@ def schedule(self) -> SchedulerOutput:
# We use `request.num_tokens` instead of
# `request.num_prompt_tokens` to consider the resumed
# requests, which have output tokens.
+
num_new_tokens = request.num_tokens - num_computed_tokens
+ if num_new_tokens <= 0:
+ # 必须保留 ≥1 个 token 重算，否则没有 logits 给采样
+ block_size = self.block_size
+ # 把 cached 长度对齐到上一个完整块边界，确保严格 < num_tokens
+ capped = ((request.num_tokens - 1) // block_size) * block_size
+ # 同步退还多出来的整块（关键！）
+ excess_blocks = (num_new_local_computed_tokens - capped) // block_size
+ if excess_blocks > 0:
+ # 具体 API 看你的 KVCacheBlocks 实现，常见的是切片或 .pop_last()
+ new_computed_blocks = new_computed_blocks[:-excess_blocks] \
+ if hasattr(new_computed_blocks, '__getitem__') \
+ else new_computed_blocks.truncate(len(new_computed_blocks) - excess_blocks)
+ num_new_local_computed_tokens = capped
+ num_computed_tokens = capped
+ num_new_tokens = request.num_tokens - num_computed_tokens
+
threshold = self.scheduler_config.long_prefill_token_threshold
if 0 < threshold < num_new_tokens:
num_new_tokens = threshold
@@ -820,7 +1130,12 @@ def schedule(self) -> SchedulerOutput:
assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens
assert token_budget >= 0
- assert len(self.running) <= self.max_num_running_reqs
+ # Fork-aware: children may overflow up to fork_child_running_overflow
+ # beyond max_num_running_reqs (Option A co-batch guarantee).
+ effective_running_cap = self.max_num_running_reqs
+ if self.enable_dynamic_fork:
+ effective_running_cap += self.fork_child_running_overflow
+ assert len(self.running) <= effective_running_cap
# Since some requests in the RUNNING queue may not be scheduled in
# this step, the total number of scheduled requests can be smaller than
# len(self.running).
@@ -871,6 +1186,12 @@ def schedule(self) -> SchedulerOutput:
self.prev_step_scheduled_req_ids.clear()
self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
+ new_block_ids_to_zero = (
+ (self.kv_cache_manager.take_new_block_ids() or None)
+ if self.needs_kv_cache_zeroing
+ else None
+ )
+
scheduler_output = SchedulerOutput(
scheduled_new_reqs=new_reqs_data,
scheduled_cached_reqs=cached_reqs_data,
@@ -886,7 +1207,13 @@ def schedule(self) -> SchedulerOutput:
# the previous and the current steps.
finished_req_ids=self.finished_req_ids,
free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
+ new_block_ids_to_zero=new_block_ids_to_zero,
+ # Fork information for Mamba state copy
+ fork_child_to_parent=self.fork_child_to_parent_this_step.copy() if self.fork_child_to_parent_this_step else None,
)
+
+ # Clear fork mapping after building scheduler output
+ self.fork_child_to_parent_this_step.clear()
# NOTE(Kuntai): this function is designed for multiple purposes:
# 1. Plan the KV cache store
@@ -912,16 +1239,111 @@ def schedule(self) -> SchedulerOutput:
def _preempt_request(self, request: Request, timestamp: float) -> None:
"""Preempt a request and put it back to the waiting queue.
+ When dynamic fork is enabled, uses partial preemption: keeps the
+ prompt's KV cache blocks allocated and only frees output blocks.
+ This avoids expensive prompt re-computation on reschedule.
+
NOTE: The request should be popped from the running queue outside of this
method.
"""
assert request.status == RequestStatus.RUNNING, (
"Only running requests can be preempted"
)
- self.kv_cache_manager.free(request)
- self.encoder_cache_manager.free(request)
+
+ # Ensure that if this request was scheduled in the previous step,
+ # we no longer treat it as such. Otherwise, when the same step both
+ # preempts and resumes the request (e.g., fork-aware admission), the
+ # resumed-path assertion in `_make_cached_request_data` will fail.
+ self.prev_step_scheduled_req_ids.discard(request.request_id)
+
+ # Partial preemption keeps the prompt-aligned KV cache prefix (and,
+ # for multimodal requests, also keeps the encoder cache reference so
+ # the GPU-side encoder outputs are not evicted), advancing
+ # num_computed_tokens to that boundary. This avoids expensive prompt
+ # re-computation on resume — including image encoder re-runs.
+ #
+ # Why partial preempt is now safe for multimodal + fork-child:
+ # - The kept KV prefix [0, tokens_kept) is a strict subset of
+ # [0, request.num_computed_tokens) at preempt time, whose KV is
+ # correctly populated.
+ # - For fork-child: the resume path goes through
+ # waiting.prepend_request and does NOT touch
+ # _process_pending_fork_requests (only fed by FORK-token emission
+ # into pending_fork_requests), so no double fork_kv_cache call /
+ # refcount corruption occurs on resume.
+ # - For multimodal: by NOT calling encoder_cache_manager.free here,
+ # the request keeps its reference on cached encoder outputs
+ # (mm_hash stays in `cached`, not in `freeable`), so they cannot
+ # be evicted by other requests' `can_allocate` while preempted.
+ # On resume, `check_and_update_cache` returns True for those
+ # mm_hashes, encoder is not rescheduled, and
+ # `_gather_mm_embeddings` finds the cached encoder outputs —
+ # avoiding the previously fatal `Encoder cache miss` assertion.
+ #
+ # Cap the keep target at num_computed_tokens to handle the rare case
+ # of preemption mid-prefill, where the request's allocated blocks may
+ # not yet cover the full num_prompt_tokens.
+ use_partial_preempt = self.enable_dynamic_fork
+
+ if use_partial_preempt:
+ # Partial preemption: keep prompt KV blocks (or as many as are
+ # actually computed) and keep encoder cache references alive.
+ block_size = self.block_size
+ tokens_to_keep = min(
+ request.num_prompt_tokens,
+ request.num_computed_tokens,
+ )
+ # Block-align down so this scheduler-side computation matches
+ # what kv_cache_manager.free_partial will actually keep.
+ tokens_to_keep = (tokens_to_keep // block_size) * block_size
+
+ # Critical: tokens_to_keep MUST NOT split a multimodal item.
+ # If [image_start, image_start+image_len) straddles
+ # tokens_to_keep (image_start < tokens_to_keep < image_end),
+ # then on resume the scheduler decides the image is not yet in
+ # KV (line ~1469), goes to check_and_update_cache, finds it in
+ # `cached` (since we don't call free here) → skips rescheduling
+ # encoding → but the worker's `encoder_cache` may have already
+ # evicted that tensor (it was made freeable when the request
+ # passed the image earlier, and another request's can_allocate
+ # could have evicted it before this preempt). The worker would
+ # then hit `assert encoder_output is not None` (Encoder cache
+ # miss). Pulling tokens_to_keep down to image_start makes the
+ # image "not started" on resume, which forces a clean re-encode
+ # path through can_allocate/allocate.
+ if request.has_encoder_inputs and request.mm_features:
+ while tokens_to_keep > 0:
+ straddler_start: int | None = None
+ for mm_feature in request.mm_features:
+ s = mm_feature.mm_position.offset
+ e = s + mm_feature.mm_position.length
+ if s < tokens_to_keep < e:
+ if straddler_start is None or s < straddler_start:
+ straddler_start = s
+ if straddler_start is None:
+ break
+ tokens_to_keep = (straddler_start // block_size) * block_size
+
+ if tokens_to_keep <= 0:
+ # Nothing safe to keep — fall back to full preempt.
+ self.kv_cache_manager.free(request)
+ request.num_computed_tokens = 0
+ self.encoder_cache_manager.free(request)
+ else:
+ tokens_kept = self.kv_cache_manager.free_partial(
+ request, tokens_to_keep)
+ request.num_computed_tokens = tokens_kept
+ # NOTE: deliberately do NOT call
+ # self.encoder_cache_manager.free here — see docstring
+ # above.
+ else:
+ # Full preemption (original behavior, used when dynamic fork is
+ # disabled). Frees both KV cache and encoder cache references.
+ self.kv_cache_manager.free(request)
+ request.num_computed_tokens = 0
+ self.encoder_cache_manager.free(request)
+
request.status = RequestStatus.PREEMPTED
- request.num_computed_tokens = 0
if request.spec_token_ids:
request.spec_token_ids = []
request.num_preemptions += 1
@@ -1352,8 +1774,10 @@ def update_from_output(
status_before_stop = request.status
# Check for stop and update request status.
+ # Also check for FORK tokens and create child requests.
+ forked = False
if new_token_ids:
- new_token_ids, stopped = self._update_request_with_output(
+ new_token_ids, stopped, forked = self._update_request_with_output(
request, new_token_ids
)
elif request.pooling_params and pooler_output is not None:
@@ -1364,6 +1788,14 @@ def update_from_output(
routed_experts = None
finish_reason = None
if stopped:
+ # Check if this request is part of a fork family
+ parent_id = self._get_fork_parent_id(req_id)
+ if parent_id is not None:
+ # This is a fork-related request
+ # Let it stop normally, but track for aggregation
+ if self.fork_debug:
+ logger.info(f"[FORK] Branch {req_id} stopped, waiting for aggregation")
+
routed_experts = self._get_routed_experts(request)
# Capture finish_reason BEFORE _handle_stopped_request, which may
@@ -1403,31 +1835,63 @@ def update_from_output(
# Get prompt logprobs for this request.
prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
+
+ # Check if this request is part of a fork family
+ parent_id = self._get_fork_parent_id(req_id)
+ is_fork_request = parent_id is not None
+
if (
new_token_ids
or pooler_output is not None
or kv_transfer_params
or stopped
):
- # Add EngineCoreOutput for this Request.
- outputs[request.client_index].append(
- EngineCoreOutput(
- request_id=req_id,
- new_token_ids=new_token_ids,
- finish_reason=finish_reason,
- new_logprobs=new_logprobs,
- new_prompt_logprobs_tensors=prompt_logprobs_tensors,
- pooling_output=pooler_output,
- stop_reason=request.stop_reason,
- events=request.take_events(),
- kv_transfer_params=kv_transfer_params,
- trace_headers=request.trace_headers,
- num_cached_tokens=request.num_cached_tokens,
- num_external_computed_tokens=request.num_external_computed_tokens,
- routed_experts=routed_experts,
- num_nans_in_logits=request.num_nans_in_logits,
+ if is_fork_request:
+ # This is a fork-related request, collect output for aggregation
+ self._collect_fork_output(parent_id, req_id, new_token_ids, stopped)
+
+ # Check if all branches are finished
+ if self._all_fork_branches_finished(parent_id):
+ # Aggregate and return the final output
+ aggregated_output = self._aggregate_fork_outputs(parent_id)
+ client_idx = self.fork_client_index[parent_id]
+ outputs[client_idx].append(
+ EngineCoreOutput(
+ request_id=parent_id, # Use parent's request_id
+ new_token_ids=aggregated_output,
+ finish_reason=finish_reason,
+ new_logprobs=None, # Logprobs not supported for aggregated output
+ new_prompt_logprobs_tensors=None,
+ pooling_output=pooler_output,
+ stop_reason=request.stop_reason,
+ events=request.take_events(),
+ kv_transfer_params=kv_transfer_params,
+ trace_headers=request.trace_headers,
+ num_cached_tokens=request.num_cached_tokens,
+ ))
+ # Clean up fork tracking data
+ self._cleanup_fork_data(parent_id)
+ else:
+ # Normal request, add output directly
+ # Add EngineCoreOutput for this Request.
+ outputs[request.client_index].append(
+ EngineCoreOutput(
+ request_id=req_id,
+ new_token_ids=new_token_ids,
+ finish_reason=finish_reason,
+ new_logprobs=new_logprobs,
+ new_prompt_logprobs_tensors=prompt_logprobs_tensors,
+ pooling_output=pooler_output,
+ stop_reason=request.stop_reason,
+ events=request.take_events(),
+ kv_transfer_params=kv_transfer_params,
+ trace_headers=request.trace_headers,
+ num_cached_tokens=request.num_cached_tokens,
+ num_external_computed_tokens=request.num_external_computed_tokens,
+ routed_experts=routed_experts,
+ num_nans_in_logits=request.num_nans_in_logits,
+ )
)
- )
else:
# Invariant: EngineCore returns no partial prefill outputs.
assert not prompt_logprobs_tensors
@@ -1458,6 +1922,9 @@ def update_from_output(
if kv_connector_output:
self._update_from_kv_xfer_finished(kv_connector_output)
+ # Process any pending fork requests created during this step
+ self._process_pending_fork_requests()
+
# collect KV cache events from KV cache manager
events = self.kv_cache_manager.take_events()
@@ -1554,21 +2021,131 @@ def _get_routed_experts(self, request: Request) -> np.ndarray | None:
def _update_request_with_output(
self, request: Request, new_token_ids: list[int]
- ) -> tuple[list[int], bool]:
+ ) -> tuple[list[int], bool, bool]:
+ """
+ Update request with new output tokens.
+
+ Returns:
+ tuple: (new_token_ids, stopped, forked)
+ - new_token_ids: The token IDs to include in output
+ - stopped: Whether the request should stop
+ - forked: Whether a fork was detected and child request created
+ """
+
# Append generated tokens and check for stop. Note that if
# a request is still being prefilled, we expect the model runner
# to return empty token ids for the request.
stopped = False
- for num_new, output_token_id in enumerate(new_token_ids, 1):
+ forked = False
+ # Convert to list if needed so we can modify it
+ new_token_ids = list(new_token_ids)
+ # Track which indices had FORK tokens (for output replacement later)
+ fork_indices = []
+
+ for idx, output_token_id in enumerate(new_token_ids):
+ # Check for FORK token before appending
+ is_fork_token = (self.enable_dynamic_fork and
+ output_token_id == self.fork_token_id)
+
+ # Prevent child branches from forking again (check if "_fork_" in request_id)
+ # If a child request tries to fork again, end its prediction early.
+ if is_fork_token and "_fork_" in request.request_id:
+ logger.info(f"[FORK] Child request {request.request_id} cannot fork again, "
+ f"ending prediction early at token index {idx}")
+ request.status = RequestStatus.FINISHED_STOPPED
+ stopped = True
+ # Drop this fork token and anything after it from the output
+ new_token_ids = new_token_ids[:idx+1]
+ fork_indices = [i for i in fork_indices if i < idx]
+ break
+
+ if is_fork_token:
+ # Debug logging for fork (only decode when debug=True)
+ if self.fork_debug:
+ output_tokens = list(request.output_token_ids) if request.output_token_ids else []
+ output_text = ""
+ if self.fork_tokenizer is not None and output_tokens:
+ try:
+ output_text = self.fork_tokenizer.decode(output_tokens, skip_special_tokens=False)
+ except Exception as e:
+ output_text = f"[decode error: {e}]"
+ is_child = getattr(request, 'is_fork_child', False)
+ logger.info(f"[FORK] request_id={request.request_id}, "
+ f"is_fork_child={is_child}, "
+ f"output_token_ids={output_tokens}, "
+ f"output_text={output_text}")
+ # time.sleep(10)
+ # Create a forked child request
+ # Note: Don't pass block_hasher=None - let it inherit from parent
+ # The parent's _block_hasher will be used for computing block_hashes
+ forked = True
+ child_request = request.fork_request(
+ child_token_id=self.child_token_id,
+ # block_hasher will be inherited from parent's _block_hasher
+ )
+ #child_request.sampling_params.max_tokens = 8000
+ #child_request.max_tokens = 8000
+
+ # Debug logging for child (only decode when debug=True)
+ if self.fork_debug:
+ child_prompt_text = ""
+ if self.fork_tokenizer is not None:
+ try:
+ child_prompt_text = self.fork_tokenizer.decode(
+ child_request.prompt_token_ids, skip_special_tokens=False
+ )
+ except Exception as e:
+ child_prompt_text = f"[decode error: {e}]"
+ logger.info(f"[FORK] Child created: request_id={child_request.request_id}, "
+ f"num_prompt_tokens={child_request.num_prompt_tokens}, "
+ f"prompt_token_ids[-20:]={child_request.prompt_token_ids[-20:] if len(child_request.prompt_token_ids) > 20 else child_request.prompt_token_ids}, "
+ f"prompt_text_suffix={child_prompt_text[-200:] if len(child_prompt_text) > 200 else child_prompt_text}")
+
+ # Queue the child request to be added after current step
+ self.pending_fork_requests.append(child_request)
+ # Track parent-child relationship
+ if request.request_id not in self.fork_relationships:
+ self.fork_relationships[request.request_id] = []
+ # Initialize fork output aggregation for this parent
+ self.fork_outputs[request.request_id] = {request.request_id: []}
+ self.fork_finished[request.request_id] = set()
+ self.fork_client_index[request.request_id] = request.client_index
+ self.fork_relationships[request.request_id].append(
+ child_request.request_id)
+ # Initialize output storage for child
+ self.fork_outputs[request.request_id][child_request.request_id] = []
+ # IMPORTANT: Keep FORK in parent's output_token_ids to maintain model context
+ # Only replace in the returned new_token_ids for display purposes
+ fork_indices.append(idx)
+
request.append_output_token_ids(output_token_id)
# Check for stop and update request state.
# This must be called before we make the EngineCoreOutput.
- stopped = check_stop(request, self.max_model_len)
- if stopped:
- del new_token_ids[num_new:] # Trim new tokens if needed.
- break
- return new_token_ids, stopped
+ # IMPORTANT: Skip stop check for FORK tokens to ensure parent continues
+ # generating more FORK tokens for other branches
+ if is_fork_token:
+ # For FORK tokens, only check length limits, not stop tokens
+ if (request.num_tokens >= self.max_model_len or
+ request.num_output_tokens >= request.max_tokens):
+ request.status = RequestStatus.FINISHED_LENGTH_CAPPED
+ stopped = True
+ new_token_ids = new_token_ids[:idx+1]
+ fork_indices = [i for i in fork_indices if i <= idx]
+ break
+ else:
+ stopped = check_stop(request, self.max_model_len)
+ if stopped:
+ new_token_ids = new_token_ids[:idx+1]
+ # Update fork_indices to match trimmed list
+ fork_indices = [i for i in fork_indices if i <= idx]
+ break
+
+ # Replace FORK with CHILD in the returned tokens (for display only)
+ for idx in fork_indices:
+ new_token_ids[idx] = self.child_token_id
+
+ return new_token_ids, stopped, forked
def _free_encoder_inputs(self, request: Request) -> None:
cached_encoder_input_ids = self.encoder_cache_manager.get_cached_input_ids(
@@ -1603,7 +2180,7 @@ def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
if request is None or request.is_finished():
# The request may have been finished. Skip.
continue
-
+
if request.is_prefill_chunk:
# Ignore draft tokens for prefill chunks.
if request.spec_token_ids:
@@ -1761,9 +2338,279 @@ def _free_request(
return kv_xfer_params
def _free_blocks(self, request: Request):
+
assert request.is_finished()
+ request_id = request.request_id
+
+ # If this is a fork child request, notify parent
+ if "_fork_" in request_id:
+ parent_id = request_id.rsplit("_fork_", 1)[0]
+ if parent_id in self.fork_relationships:
+ child_list = self.fork_relationships[parent_id]
+ if request_id in child_list:
+ child_list.remove(request_id)
+ if self.fork_debug:
+ logger.info(f"[Fork] Removed child {request_id} from parent's fork_relationships")
+
self.kv_cache_manager.free(request)
- del self.requests[request.request_id]
+ del self.requests[request_id]
+
+ # =========================================================================
+ # Fork-related helper methods
+ # =========================================================================
+
+ def _get_fork_parent_id(self, req_id: str) -> str | None:
+ """Get the root parent request ID if this request is part of a fork family.
+
+ For nested forks like chatcmpl-xxx_fork_1_fork_1, this returns chatcmpl-xxx
+ (the original parent), not chatcmpl-xxx_fork_1.
+ """
+ # Check if this is a parent request itself
+ if req_id in self.fork_outputs:
+ return req_id
+
+ # Check if this is a child request (format: parent_id_fork_N or nested)
+ if "_fork_" not in req_id:
+ return None
+
+ # For nested forks, we need to find the root parent
+ # chatcmpl-xxx_fork_1_fork_1 -> chatcmpl-xxx (split on first _fork_)
+ root_parent_id = req_id.split("_fork_", 1)[0]
+ if root_parent_id in self.fork_outputs:
+ return root_parent_id
+
+ # Also check immediate parent (for compatibility with existing nested fork tracking)
+ immediate_parent_id = req_id.rsplit("_fork_", 1)[0]
+ if immediate_parent_id in self.fork_outputs:
+ return immediate_parent_id
+
+ return None
+
+ def _collect_fork_output(self, parent_id: str, req_id: str,
+ new_token_ids: list[int], stopped: bool) -> None:
+ """Collect output tokens from a fork branch."""
+ if parent_id not in self.fork_outputs:
+ return
+
+ # Ensure this request has an output list
+ if req_id not in self.fork_outputs[parent_id]:
+ self.fork_outputs[parent_id][req_id] = []
+
+ # Append new tokens
+ self.fork_outputs[parent_id][req_id].extend(new_token_ids)
+
+ # Mark as finished if stopped
+ if stopped:
+ self.fork_finished[parent_id].add(req_id)
+ # Family statistics changed — remaining siblings' priority may shift
+ self._fork_queue_dirty = True
+ if self.fork_debug:
+ logger.info(f"[FORK] Branch {req_id} finished, "
+ f"total tokens: {len(self.fork_outputs[parent_id][req_id])}")
+
+ def _all_fork_branches_finished(self, parent_id: str) -> bool:
+ """Check if all fork branches (parent + children) have finished."""
+ if parent_id not in self.fork_outputs:
+ return False
+
+ # Get all branch IDs (parent + children)
+ all_branch_ids = set(self.fork_outputs[parent_id].keys())
+ finished_ids = self.fork_finished.get(parent_id, set())
+
+ all_finished = all_branch_ids == finished_ids
+ if all_finished:
+ logger.info(f"[FORK] All {len(all_branch_ids)} branches finished for parent {parent_id}")
+
+ return all_finished
+
+ def _aggregate_fork_outputs(self, parent_id: str) -> list[int]:
+ """Aggregate outputs from all fork branches into a single output.
+
+ The aggregation logic:
+ 1. Iterate through parent's tokens
+ 2. When encountering a CHILD token, insert the corresponding child branch's output after it
+ 3. Child branches are matched by their fork number (_fork_1, _fork_2, ...)
+
+ Result format: ...[child_1_content]...[child_2_content]...
+ """
+
+ if parent_id not in self.fork_outputs:
+ return []
+
+ branch_outputs = self.fork_outputs[parent_id]
+ parent_tokens = branch_outputs.get(parent_id, [])
+
+ # Get child branches sorted by fork number
+ # Format: parent_id_fork_1, parent_id_fork_2, ...
+ child_branches = [(bid, branch_outputs[bid])
+ for bid in branch_outputs.keys()
+ if bid != parent_id]
+ # Sort by fork number (extract the number after "_fork_")
+ child_branches.sort(key=lambda x: int(x[0].rsplit("_fork_", 1)[1]))
+
+ logger.info(f"[FORK] Aggregating: parent has {len(parent_tokens)} tokens, "
+ f"{len(child_branches)} child branches")
+
+ # Interleave parent tokens with child outputs
+ aggregated: list[int] = []
+ child_idx = 0
+
+ for token in parent_tokens:
+ aggregated.append(token)
+ # When we see a CHILD token, insert the corresponding child's output after it
+ if token == self.child_token_id and child_idx < len(child_branches):
+ child_id, child_tokens = child_branches[child_idx]
+ aggregated.extend(child_tokens)
+ if self.fork_debug:
+ logger.info(f"[FORK] Inserted child {child_idx + 1} ({child_id}): "
+ f"{len(child_tokens)} tokens after CHILD token")
+ child_idx += 1
+
+ # Log warning if not all children were inserted
+ if child_idx < len(child_branches):
+ if self.fork_debug:
+ logger.warning(f"[FORK] Not all children inserted! "
+ f"Expected {len(child_branches)}, inserted {child_idx}")
+ if self.fork_debug:
+ logger.info(f"[FORK] Total aggregated output: {len(aggregated)} tokens")
+ return aggregated
+
+ def _cleanup_fork_data(self, parent_id: str) -> None:
+ """Clean up fork tracking data after aggregation.
+
+ Also handles nested fork cleanup - if parent_id itself contains '_fork_',
+ it means this is a nested fork family that also needs cleanup.
+ """
+
+ # Clean up the direct fork data
+ if parent_id in self.fork_outputs:
+ del self.fork_outputs[parent_id]
+ if parent_id in self.fork_finished:
+ del self.fork_finished[parent_id]
+ if parent_id in self.fork_client_index:
+ del self.fork_client_index[parent_id]
+ if parent_id in self.fork_relationships:
+ del self.fork_relationships[parent_id]
+
+ # Also clean up any nested fork data that may have been created
+ # by child requests that also generated FORK tokens
+ nested_ids_to_clean = []
+ for fork_parent in list(self.fork_outputs.keys()):
+ if fork_parent.startswith(parent_id + "_fork_"):
+ nested_ids_to_clean.append(fork_parent)
+
+ for nested_id in nested_ids_to_clean:
+ if nested_id in self.fork_outputs:
+ del self.fork_outputs[nested_id]
+ if nested_id in self.fork_finished:
+ del self.fork_finished[nested_id]
+ if nested_id in self.fork_client_index:
+ del self.fork_client_index[nested_id]
+ if nested_id in self.fork_relationships:
+ del self.fork_relationships[nested_id]
+ if self.fork_debug:
+ logger.info(f"[FORK] Cleaned up nested fork data for {nested_id}")
+ if self.fork_debug:
+ logger.info(f"[FORK] Cleaned up fork data for parent {parent_id}")
+
+ def _process_pending_fork_requests(self) -> None:
+ """Add pending fork requests to the scheduler."""
+ if not self.pending_fork_requests:
+ return
+
+ if self.fork_debug:
+ logger.info(f"[FORK] Processing {len(self.pending_fork_requests)} pending fork requests")
+
+ # Get block_size for calculating num_computed_tokens
+ # IMPORTANT: Use the actual KV cache block size from kv_cache_config, NOT self.block_size
+ # self.block_size is max_model_len (e.g., 16000), but the actual KV cache block size
+ # is from kv_cache_spec (e.g., 544). Using the wrong block_size causes:
+ # - num_full_blocks = fork_position // 16000 = 0 (no blocks shared)
+ # - Child branches don't inherit any KV cache, causing garbled output
+ kv_cache_groups = self.kv_cache_config.kv_cache_groups
+ if kv_cache_groups:
+ block_size = kv_cache_groups[0].kv_cache_spec.block_size
+ else:
+ block_size = self.block_size
+
+ for child_request in self.pending_fork_requests:
+ # DEBUG: Print detailed info
+ if self.fork_debug:
+ print(f"[FORK PROCESS DEBUG] child_request.request_id={child_request.request_id}")
+ print(f"[FORK PROCESS DEBUG] child_request.num_prompt_tokens={child_request.num_prompt_tokens}")
+ print(f"[FORK PROCESS DEBUG] child_request.num_output_tokens={child_request.num_output_tokens}")
+ print(f"[FORK PROCESS DEBUG] child_request.num_tokens={child_request.num_tokens}")
+ print(f"[FORK PROCESS DEBUG] len(prompt_token_ids)={len(child_request.prompt_token_ids) if child_request.prompt_token_ids else 0}")
+ print(f"[FORK PROCESS DEBUG] block_size={block_size}")
+
+ # In method B: child inherits parent's output_token_ids
+ # fork_position = child's total tokens = prompt + inherited output (including CHILD token)
+ # This is the total sequence length that needs KV cache
+ fork_position = child_request.num_tokens # = num_prompt_tokens + num_output_tokens
+
+ # Initialize num_shared_tokens to 0 (will be updated based on actual shared blocks)
+ num_shared_tokens = 0
+
+ if child_request.parent_request_id:
+ # Try to fork KV cache (share parent's FULL blocks only)
+ try:
+ shared_blocks = self.kv_cache_manager.fork_kv_cache(
+ parent_request_id=child_request.parent_request_id,
+ child_request_id=child_request.request_id,
+ fork_position=fork_position - 1, # Position before CHILD token
+ block_size=block_size, # Pass scheduler's block_size
+ )
+ # IMPORTANT: For hybrid models (Attention + Mamba), we need to find
+ # the first non-empty block group. blocks[0] might be MambaManager's
+ # empty list, while blocks[1] is AttentionManager's shared blocks.
+ num_shared_blocks = 0
+ if shared_blocks.blocks:
+ for block_group in shared_blocks.blocks:
+ if block_group: # Find first non-empty block group
+ num_shared_blocks = len(block_group)
+ break
+ # IMPORTANT: Calculate num_shared_tokens based on ACTUAL shared blocks
+ # Use the block_size from kv_cache_manager, not scheduler
+ if num_shared_blocks > 0:
+ num_shared_tokens = num_shared_blocks * block_size
+ else:
+ num_shared_tokens = 0 # No blocks shared, recompute everything
+ if self.fork_debug:
+ logger.info(f"[FORK] Shared {num_shared_blocks} full blocks from parent "
+ f"({num_shared_tokens} tokens are shared)")
+ except Exception as e:
+ logger.warning(f"Failed to fork KV cache: {e}")
+ num_shared_tokens = 0 # Fallback: recompute everything
+
+ # Set num_computed_tokens = number of tokens covered by shared FULL blocks
+ # Child will recompute tokens from num_shared_tokens to fork_position
+ # This includes the partial block tokens AND the CHILD token
+ #child_request.num_computed_tokens = num_shared_tokens
+ child_request.num_computed_tokens = min(num_shared_tokens, child_request.num_tokens - 1)
+
+ if self.fork_debug:
+ logger.info(f"[FORK] Child {child_request.request_id}: "
+ f"num_prompt_tokens={child_request.num_prompt_tokens}, "
+ f"num_output_tokens={child_request.num_output_tokens}, "
+ f"fork_position={fork_position}, "
+ f"num_shared_tokens={num_shared_tokens}, "
+ f"num_computed_tokens={child_request.num_computed_tokens}, "
+ f"tokens_to_recompute={fork_position - num_shared_tokens}")
+
+ # Add the child request to the scheduler
+ self.add_request(child_request)
+ if self.fork_debug:
+ logger.info(f"[FORK] Added child request {child_request.request_id} to scheduler")
+
+ # Track fork child->parent mapping for Mamba state copy
+ # This will be used by GPU model runner to copy Mamba state from parent
+ if child_request.parent_request_id:
+ self.fork_child_to_parent_this_step[child_request.request_id] = child_request.parent_request_id
+
+ # Mark waiting queue dirty so next schedule() re-sorts with new children
+ self._fork_queue_dirty = True
+ # Clear the pending list
+ self.pending_fork_requests.clear()
@property
def pause_state(self) -> PauseState:
diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py
index c071ae155c94..00dca76de162 100644
--- a/vllm/v1/core/single_type_kv_cache_manager.py
+++ b/vllm/v1/core/single_type_kv_cache_manager.py
@@ -55,6 +55,7 @@ def __init__(
self.kv_cache_spec = kv_cache_spec
self.block_pool = block_pool
self.enable_caching = enable_caching
+ self.new_block_ids: list[int] = []
# Mapping from request ID to blocks to track the blocks allocated
# for each request, so that we can free the blocks when the request
@@ -208,6 +209,8 @@ def allocate_new_computed_blocks(
cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks)
)
req_blocks.extend(allocated_blocks)
+ if type(self.kv_cache_spec) is FullAttentionSpec:
+ self.new_block_ids.extend(b.block_id for b in allocated_blocks)
def allocate_new_blocks(
self, request_id: str, num_tokens: int, num_tokens_main_model: int
@@ -234,8 +237,16 @@ def allocate_new_blocks(
else:
new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
req_blocks.extend(new_blocks)
+ if type(self.kv_cache_spec) is FullAttentionSpec:
+ self.new_block_ids.extend(b.block_id for b in new_blocks)
return new_blocks
+ def take_new_block_ids(self) -> list[int]:
+ """Drain and return block IDs allocated since the last call."""
+ ids = self.new_block_ids
+ self.new_block_ids = []
+ return ids
+
def cache_blocks(self, request: Request, num_tokens: int) -> None:
"""
Cache the blocks for the request.
@@ -279,6 +290,47 @@ def free(self, request_id: str) -> None:
self.block_pool.free_blocks(ordered_blocks)
self.num_cached_block.pop(request_id, None)
+ def free_tail_blocks(self, request_id: str,
+ num_blocks_to_keep: int) -> int:
+ """Free tail blocks for a request, keeping the first
+ num_blocks_to_keep blocks allocated.
+
+ Args:
+ request_id: The request ID.
+ num_blocks_to_keep: Number of prefix blocks to retain.
+
+ Returns:
+ The number of blocks actually kept.
+ """
+ req_blocks = self.req_to_blocks.get(request_id, [])
+ if not req_blocks:
+ return 0
+
+ # Clamp to actual block count
+ num_blocks_to_keep = max(0, min(num_blocks_to_keep, len(req_blocks)))
+ if num_blocks_to_keep >= len(req_blocks):
+ # Nothing to free
+ return len(req_blocks)
+
+ # Free tail blocks in reverse order (tail-first for LRU ordering)
+ blocks_to_free = req_blocks[num_blocks_to_keep:]
+ self.block_pool.free_blocks(reversed(blocks_to_free))
+
+ # Keep only prefix blocks
+ self.req_to_blocks[request_id] = req_blocks[:num_blocks_to_keep]
+ # Update num_cached_block to reflect the kept blocks that already
+ # have block_hash set. Clearing it entirely causes cache_blocks()
+ # to re-cache blocks from position 0, hitting the assertion
+ # `assert blk.block_hash is None` for blocks that already have hashes.
+ old_cached = self.num_cached_block.get(request_id, 0)
+ if num_blocks_to_keep > 0:
+ self.num_cached_block[request_id] = min(
+ num_blocks_to_keep, old_cached)
+ else:
+ self.num_cached_block.pop(request_id, None)
+
+ return num_blocks_to_keep
+
@abstractmethod
def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
"""
diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py
index ad70f839d9c5..8b7d9f6abc63 100644
--- a/vllm/v1/engine/input_processor.py
+++ b/vllm/v1/engine/input_processor.py
@@ -227,7 +227,7 @@ def process_inputs(
) -> EngineCoreRequest:
self._validate_params(params, supported_tasks)
self._validate_lora(lora_request)
-
+
parallel_config = self.vllm_config.parallel_config
dp_size = parallel_config.data_parallel_size
dp_local_size = parallel_config.data_parallel_size_local
@@ -248,7 +248,13 @@ def process_inputs(
if arrival_time is None:
arrival_time = prompt.get("arrival_time", time.time()) # type: ignore[assignment]
-
+ # <|im_start|>user\ndocument parsing.\n<|im_end|>\n<|im_start|>assistant\ntext [36, 77, 272, 134]<|im_end|>\n<|im_start|>assistant\n
+ # 151644 (<|im_start|>) 151645 (<|im_end|>)
+ if prompt['prompt_token_ids'].count(151645) >= 2:
+ prompt['prompt_token_ids'] = prompt['prompt_token_ids'][:-5]
+ elif prompt['prompt_token_ids'].count(248046) >= 2:
+ prompt['prompt_token_ids'] = prompt['prompt_token_ids'][:-9]
+ # print(prompt)
processed_inputs: ProcessorInputs = prompt # type: ignore[assignment]
else:
logger.warning_once(
@@ -307,6 +313,24 @@ def process_inputs(
decoder_mm_positions = decoder_inputs["mm_placeholders"]
decoder_mm_hashes = decoder_inputs["mm_hashes"]
+ # Print image token count
+ '''prompt_len = len(prompt_token_ids) if prompt_token_ids else 0
+ for modality, mm_positions in decoder_mm_positions.items():
+ num_images = len(mm_positions)
+ image_tokens = sum(p.get_num_embeds() for p in mm_positions)
+ tokens_per_tile = 256 # (image_size/patch_size)^2 * downsample_ratio^2
+ num_tiles = image_tokens // tokens_per_tile if tokens_per_tile else 0
+ text_tokens = prompt_len - image_tokens
+ logger.info(
+ "request_id=%s modality=%s images=%d tiles=%d "
+ "image_tokens=%d text_tokens=%d prompt_total=%d "
+ "image_ratio=%.1f%%",
+ request_id, modality, num_images, num_tiles,
+ image_tokens, text_tokens, prompt_len,
+ image_tokens / prompt_len * 100 if prompt_len else 0,
+ )
+ '''
+
if not all(
isinstance(leaf, str) for leaf in json_iter_leaves(decoder_mm_hashes)
):
diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py
index 4a1b16fc580c..48ecf6b9dc85 100644
--- a/vllm/v1/kv_cache_interface.py
+++ b/vllm/v1/kv_cache_interface.py
@@ -489,3 +489,11 @@ class KVCacheConfig:
For models with multiple types of attention, there will be multiple groups,
see `_get_kv_cache_config_uniform_page_size` for more details.
"""
+
+ @property
+ def has_mamba_layers(self) -> bool:
+ return any(isinstance(g.kv_cache_spec, MambaSpec) for g in self.kv_cache_groups)
+
+ @property
+ def needs_kv_cache_zeroing(self) -> bool:
+ return self.has_mamba_layers
diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py
index 22b06f0e2d97..1dcf0afbac3f 100644
--- a/vllm/v1/outputs.py
+++ b/vllm/v1/outputs.py
@@ -206,6 +206,21 @@ class ECConnectorOutput:
finished_recving: set[str] | None = None
+@dataclass
+class ForkInfo:
+ """Information about a fork request triggered by FORK_TOKEN detection."""
+ # The parent request ID that triggered the fork
+ parent_request_id: str
+ # The token position where fork occurred
+ fork_position: int
+ # The prefix token IDs for the child request (prompt + branch content)
+ prefix_token_ids: list[int]
+ # Index of this fork (for generating unique child request ID)
+ fork_idx: int
+ # The block start position in parent for KV cache sharing
+ block_start_position: int = 0
+
+
# ModelRunnerOutput is serialized and sent to the scheduler process.
# This is expensive for torch.Tensor so prefer to use list instead.
@dataclass
@@ -247,6 +262,10 @@ class ModelRunnerOutput:
# information related to cudagraph execution
cudagraph_stats: CUDAGraphStat | None = None
+ # Fork information for dynamic fork inference
+ # Contains info about requests that need to fork due to FORK_TOKEN detection
+ fork_infos: list[ForkInfo] | None = None
+
# ModelRunnerOutput wrapper for async scheduling.
class AsyncModelRunnerOutput(ABC):
diff --git a/vllm/v1/request.py b/vllm/v1/request.py
index 85ca90d9957c..102a93858703 100644
--- a/vllm/v1/request.py
+++ b/vllm/v1/request.py
@@ -74,6 +74,9 @@ def __init__(
block_hasher: Callable[["Request"], list["BlockHash"]] | None = None,
resumable: bool = False,
reasoning_ended: bool | None = None,
+ # Fork-related fields
+ parent_request_id: str | None = None,
+ is_fork_child: bool = False,
) -> None:
self.request_id = request_id
self.client_index = client_index
@@ -92,6 +95,12 @@ def __init__(
self.events: list[EngineCoreEvent] = []
self.stop_reason: int | str | None = None
+ # Fork-related fields
+ self.parent_request_id = parent_request_id
+ self.is_fork_child = is_fork_child
+ self.child_request_ids: list[str] = [] # Track forked children
+ self.fork_count = 0 # Counter for generating unique child IDs
+
# P/D: Connector-specific KV transfer parameters.
self.kv_transfer_params: dict[str, Any] | None = None
@@ -303,6 +312,105 @@ def __lt__(self, other: "Request") -> bool:
return self.request_id < other.request_id
return id(self) < id(other)
+ def fork_request(
+ self,
+ child_token_id: int,
+ block_hasher=None,
+ ) -> "Request":
+ """
+ Create a forked child request from the current request.
+
+ The child request inherits:
+ - All prompt tokens + output tokens so far
+ - Replaces the last token (FORK) with child_token_id (CHILD)
+ - Same sampling params, mm_features, etc.
+
+ Args:
+ child_token_id: The token ID to use instead of FORK (typically CHILD token)
+ block_hasher: Optional block hasher for prefix caching.
+ If None, will inherit from parent's _block_hasher.
+
+ Returns:
+ A new Request object representing the forked child
+ """
+ import copy
+
+ debug = False
+ self.fork_count += 1
+ child_request_id = f"{self.request_id}_fork_{self.fork_count}"
+
+ # Build child's token sequence:
+ # prompt_tokens + output_tokens (before FORK) + [child_token_id]
+ # Note: FORK token is NOT in _output_token_ids yet when this is called
+ child_prompt_tokens = self.prompt_token_ids.copy()
+
+ # DEBUG: Log the construction of child prompt
+ if debug:
+ print(f"[FORK DEBUG] Parent request_id={self.request_id}")
+ print(f"[FORK DEBUG] Parent prompt_token_ids length={len(self.prompt_token_ids)}")
+ print(f"[FORK DEBUG] Parent _output_token_ids length={len(self._output_token_ids)}")
+ print(f"[FORK DEBUG] Parent _output_token_ids={self._output_token_ids}")
+ print(f"[FORK DEBUG] child_token_id={child_token_id}")
+
+ if self._output_token_ids:
+ # Include all current output tokens (FORK not included yet)
+ # These are tokens BEFORE the FORK was generated
+ child_prompt_tokens.extend(self._output_token_ids)
+ # Add CHILD token as the new last token (this replaces FORK)
+ child_prompt_tokens.append(child_token_id)
+ if debug:
+ print(f"[FORK DEBUG] Final child_prompt_tokens length={len(child_prompt_tokens)}")
+ print(f"[FORK DEBUG] Final child_prompt_tokens[-10:]={child_prompt_tokens[-10:]}")
+
+ # Deep copy sampling params to avoid sharing state
+ child_sampling_params = copy.deepcopy(self.sampling_params) if self.sampling_params else None
+
+ # IMPORTANT: Inherit parent's block_hasher if not specified
+ # This is required for prefix caching to work correctly.
+ # Without block_hasher, block_hashes will be empty, causing:
+ # AssertionError: len(request.block_hashes) >= num_full_blocks
+ effective_block_hasher = block_hasher if block_hasher is not None else self._block_hasher
+
+ child_request = Request(
+ request_id=child_request_id,
+ prompt_token_ids=child_prompt_tokens,
+ sampling_params=child_sampling_params,
+ pooling_params=self.pooling_params,
+ client_index=self.client_index,
+ arrival_time=time.time(),
+ mm_features=self.mm_features, # Share mm_features (read-only)
+ lora_request=self.lora_request,
+ cache_salt=self.cache_salt,
+ priority=self.priority,
+ trace_headers=self.trace_headers,
+ block_hasher=effective_block_hasher,
+ parent_request_id=self.request_id,
+ is_fork_child=True,
+ )
+
+ # NOTE: Do NOT copy parent's block_hashes to child!
+ #
+ # The block_hasher function uses len(block_hashes) * block_size as the
+ # start_token_idx to determine which blocks need to be hashed.
+ # If we copy parent's block_hashes, the hasher will think those blocks
+ # are already hashed and skip them. But when child generates new tokens
+ # and the scheduler calls cache_blocks(), it expects block_hashes to
+ # include all full blocks up to num_computed_tokens.
+ #
+ # By NOT copying block_hashes, the Request.__init__ will call
+ # update_block_hashes() which computes block_hashes for the child's
+ # entire token sequence (prompt_tokens + inherited output_tokens).
+ # This ensures block_hashes is always consistent with the child's tokens.
+ #
+ # The child still shares parent's KV cache blocks via fork_kv_cache(),
+ # but block_hashes are computed independently based on token content.
+
+ # Track the child in parent
+ self.child_request_ids.append(child_request_id)
+
+ return child_request
+
+
class RequestStatus(enum.IntEnum):
"""Status of a request."""
diff --git a/vllm/v1/spec_decode/medusa.py b/vllm/v1/spec_decode/medusa.py
index 80b0f0a9870a..ce1ae8c131f8 100644
--- a/vllm/v1/spec_decode/medusa.py
+++ b/vllm/v1/spec_decode/medusa.py
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import time
import torch
import torch.nn as nn
@@ -39,21 +40,30 @@ def __init__(
def propose(
self,
target_hidden_states: torch.Tensor,
+ input_ids: torch.Tensor,
+ inputs_embeds: torch.Tensor,
sampling_metadata: SamplingMetadata,
slot_mappings: dict[str, torch.Tensor]
| list[dict[str, torch.Tensor]]
| None = None, # unused
) -> torch.Tensor:
# Generate blocks and compute logits
- blocks = self.model(target_hidden_states)
- logits = self.model.compute_logits(blocks)
+ # blocks = self.model(hidden_states=target_hidden_states, inputs_embeds=inputs_embeds)
+ # logits = self.model.compute_logits(blocks)
# Compute argmax for each Medusa head and stack into a single tensor
# Shape: [batch_size, num_heads]
- draft_tokens = torch.stack([logit.argmax(dim=-1) for logit in logits], dim=1)
+ # draft_tokens = torch.stack([logit.argmax(dim=-1) for logit in logits], dim=1)
+ # t0 = time.time()
+ draft_tokens = self.model(hidden_states=target_hidden_states, inputs_embeds=inputs_embeds, input_ids=input_ids)
+ # print('proposed times:', time.time() - t0)
return draft_tokens
+ # draft_tokens = self.model(hidden_states=target_hidden_states, inputs_embeds=inputs_embeds)
+ # draft_tokens = [logit.argmax(dim=-1).tolist() for logit in logits]
+ # return [list(row) for row in zip(*draft_tokens)]
+
def load_model(self, target_model: nn.Module) -> None:
from vllm.compilation.backends import set_model_tag
@@ -69,10 +79,12 @@ def load_model(self, target_model: nn.Module) -> None:
@torch.inference_mode()
def dummy_run(self, num_tokens: int) -> None:
- hidden_states = torch.zeros(
- (self.max_num_tokens, self.hidden_size),
- dtype=self.dtype,
- device=self.device,
- )
- with set_forward_context(None, self.vllm_config, num_tokens=num_tokens):
- self.model(hidden_states)
+ hidden_states = torch.zeros((self.max_num_tokens, self.hidden_size),
+ dtype=self.dtype,
+ device=self.device)
+ inputs_embeds = torch.zeros((self.max_num_tokens, self.hidden_size),
+ dtype=self.dtype,
+ device=self.device)
+ with set_forward_context(None, self.vllm_config,
+ num_tokens=num_tokens):
+ self.model(hidden_states, inputs_embeds)
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
index e4ddefc8180f..fdead97dea90 100644
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -187,6 +187,7 @@
from .utils import (
AttentionGroup,
+ KVBlockZeroer,
add_kv_sharing_layers_to_kv_cache_groups,
bind_kv_cache,
prepare_kernel_block_sizes,
@@ -918,6 +919,26 @@ def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:
decode_threshold=self.reorder_batch_threshold,
)
+ def _init_kv_zero_meta(self) -> None:
+ """One-time precomputation for _zero_block_ids.
+
+ Delegates to KVBlockZeroer.init_meta with the runner's state.
+ Called from gpu_worker.py outside the CuMem pool context.
+ """
+ self._kv_block_zeroer = KVBlockZeroer(self.device, self.pin_memory)
+ self._kv_block_zeroer.init_meta(
+ attn_groups_iter=self._kv_cache_spec_attn_group_iterator(),
+ kernel_block_sizes=self._kernel_block_sizes,
+ cache_dtype=self.cache_config.cache_dtype,
+ runner_only_attn_layers=self.runner_only_attn_layers,
+ static_forward_context=(self.compilation_config.static_forward_context),
+ )
+
+ def _zero_block_ids(self, block_ids: list[int]) -> None:
+ """Zero the KV cache memory for the given block IDs."""
+ if hasattr(self, "_kv_block_zeroer"):
+ self._kv_block_zeroer.zero_block_ids(block_ids)
+
# Note: used for model runner override.
def _init_device_properties(self) -> None:
"""Initialize attributes from torch.cuda.get_device_properties"""
@@ -951,6 +972,11 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
for req_id in scheduler_output.finished_req_ids:
self.input_batch.remove_request(req_id)
+ # Zero GPU memory for freshly allocated cache blocks to prevent
+ # stale NaN/data from corrupting attention or SSM computation.
+ if scheduler_output.new_block_ids_to_zero:
+ self._zero_block_ids(scheduler_output.new_block_ids_to_zero)
+
# Free the cached encoder outputs.
for mm_hash in scheduler_output.free_encoder_mm_hashes:
self.encoder_cache.pop(mm_hash, None)
@@ -3795,7 +3821,7 @@ def propose_draft_token_ids(sampled_token_ids):
if input_fits_in_drafter:
propose_draft_token_ids(sampled_token_ids)
elif self.valid_sampled_token_count_event is not None:
- assert spec_decode_common_attn_metadata is not None
+ assert spec_decode_common_attn_metadata is not None
next_token_ids, valid_sampled_tokens_count = (
self.drafter.prepare_next_token_ids_padded(
spec_decode_common_attn_metadata,
@@ -4073,9 +4099,13 @@ def propose_draft_token_ids(
offset += num_draft + 1
indices = torch.tensor(indices, device=self.device)
hidden_states = sample_hidden_states[indices]
+ block_token = 151678 #248077
+ input_ids_for_draft = torch.tensor([sampled_token_id[-1] if len(sampled_token_id) else block_token for sampled_token_id in sampled_token_ids ], device=self.device)
draft_token_ids = self.drafter.propose(
target_hidden_states=hidden_states,
+ inputs_embeds=None,
+ input_ids=input_ids_for_draft,
sampling_metadata=sampling_metadata,
slot_mappings=slot_mappings,
)
@@ -5554,7 +5584,7 @@ def initialize_metadata_builders(
or self.speculative_config.uses_draft_model()
):
assert isinstance(self.drafter, EagleProposer | DraftModelProposer)
- self.drafter.initialize_attn_backend(kv_cache_config, kernel_block_sizes)
+ self.drafter.initialize_attn_backend(kv_cache_config, kernel_block_sizes)
def _check_and_update_cudagraph_mode(
self,
@@ -6066,6 +6096,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
kernel_block_sizes = prepare_kernel_block_sizes(
kv_cache_config, self.attn_groups
)
+ self._kernel_block_sizes = kernel_block_sizes
# create metadata builders
self.initialize_metadata_builders(kv_cache_config, kernel_block_sizes)
diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py
index c0654abd53a2..0b74d450ad50 100644
--- a/vllm/v1/worker/gpu_worker.py
+++ b/vllm/v1/worker/gpu_worker.py
@@ -480,6 +480,14 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
else:
self.model_runner.initialize_kv_cache(kv_cache_config)
+ # Build KV-zero metadata outside the CuMem pool so the bookkeeping
+ # GPU tensors (seg_addrs, block-id buffers) use the standard PyTorch
+ # allocator and are not discarded during sleep/wake cycles.
+ if kv_cache_config.needs_kv_cache_zeroing and hasattr(
+ self.model_runner, "_init_kv_zero_meta"
+ ):
+ self.model_runner._init_kv_zero_meta()
+
@instrument(span_name="Warmup (GPU)")
def compile_or_warm_up_model(self) -> float:
warmup_sizes = []
diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py
index 2bd5d2b3fea8..8840d5fab9b2 100644
--- a/vllm/v1/worker/mamba_utils.py
+++ b/vllm/v1/worker/mamba_utils.py
@@ -140,6 +140,20 @@ def do_mamba_copy_block(copy_bufs: MambaCopyBuffers):
)
+def _get_parent_req_id_from_fork_id(req_id: str) -> str | None:
+ """
+ Extract parent request ID from fork child request ID.
+ Fork child request ID format: parent_id_fork_N (e.g., "chatcmpl-xxx_fork_1")
+
+ FIX for Issue 1: This allows us to get parent_req_id even if fork_child_to_parent
+ mapping was cleared before the child was scheduled.
+ """
+ if "_fork_" not in req_id:
+ return None
+ # Get immediate parent (for nested forks like parent_fork_1_fork_2, returns parent_fork_1)
+ return req_id.rsplit("_fork_", 1)[0]
+
+
def preprocess_mamba(
scheduler_output: SchedulerOutput,
kv_cache_config: KVCacheConfig,
@@ -154,6 +168,9 @@ def preprocess_mamba(
"""
Copy the mamba state of previous step to the last
(1 + num_speculative_blocks) block.
+
+ For forked child requests, copy Mamba state from parent request instead
+ of using the child's own (uninitialized) state.
"""
mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config)
num_speculative_blocks = mamba_spec.num_speculative_blocks
@@ -171,10 +188,71 @@ def preprocess_mamba(
for req_id in itertools.chain(finished_req_ids, preempted_req_ids, resumed_req_ids):
mamba_state_idx.pop(req_id, None)
+ # Get fork information: child_req_id -> parent_req_id
+ fork_child_to_parent = scheduler_output.fork_child_to_parent or {}
+
copy_bufs.offset = 0
for i, req_id in enumerate(input_batch.req_ids):
req_state = requests[req_id]
prev_state_idx = mamba_state_idx.get(req_id)
+
+ # Check if this is a forked child request that needs to copy from parent
+ # FIX Issue 1: Try fork_child_to_parent first, then fallback to parsing request_id
+ # This handles the case where fork_child_to_parent mapping was cleared before
+ # the child was actually scheduled (delayed scheduling).
+ parent_req_id = fork_child_to_parent.get(req_id)
+ if parent_req_id is None and prev_state_idx is None:
+ # Fallback: try to extract parent_req_id from request_id string
+ parent_req_id = _get_parent_req_id_from_fork_id(req_id)
+
+ if parent_req_id is not None and prev_state_idx is None:
+ # This is a NEW fork child request - need to copy from parent's Mamba state
+ if parent_req_id in requests:
+ parent_req_state = requests[parent_req_id]
+
+ # FIX Issue 2: Calculate CORRECT source state position
+ # Child starts from num_computed_tokens, so we need the Mamba state
+ # at that position, NOT the parent's current position.
+ #
+ # Example: parent at token 100, child num_computed_tokens = 96 (shared blocks)
+ # - Child needs state from block covering position 95 (last token before 96)
+ # - This is block index = (96 - 1) // block_size = 95 // block_size
+ # - NOT parent's current state which might be at token 100
+ child_start_position = req_state.num_computed_tokens
+ if child_start_position > 0:
+ # The correct source state is the block that contains position (child_start_position - 1)
+ # This is the Mamba state BEFORE the child starts computing
+ correct_source_state_idx = (child_start_position - 1) // block_size
+ else:
+ correct_source_state_idx = -1
+
+ # Calculate child's destination state index
+ num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
+ child_num_blocks: int = (
+ cdiv(req_state.num_computed_tokens + num_scheduled_tokens, block_size)
+ + num_speculative_blocks
+ )
+ child_state_idx = child_num_blocks - 1 - num_speculative_blocks
+ mamba_state_idx[req_id] = child_state_idx
+
+ # Only copy if we have a valid source state
+ if correct_source_state_idx >= 0:
+ # Collect copy metadata from parent to child
+ # Use CORRECT source position, not parent's current position
+ collect_mamba_copy_meta_cross_request(
+ copy_bufs,
+ kv_cache_config,
+ mamba_state_copy_funcs,
+ mamba_group_ids,
+ correct_source_state_idx, # FIX: Use correct source position
+ child_state_idx,
+ parent_req_state,
+ req_state,
+ forward_context,
+ )
+ input_batch.num_accepted_tokens_cpu[i] = 1
+ continue
+
if prev_state_idx is None:
# new / resumed request, no previous state
# if num_computed_tokens is 0, prev_state_idx will be -1
@@ -214,6 +292,48 @@ def preprocess_mamba(
do_mamba_copy_block(copy_bufs)
+def collect_mamba_copy_meta_cross_request(
+ copy_bufs: MambaCopyBuffers,
+ kv_cache_config: KVCacheConfig,
+ mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
+ mamba_group_ids: list[int],
+ src_block_idx: int,
+ dest_block_idx: int,
+ src_req_state: CachedRequestState,
+ dest_req_state: CachedRequestState,
+ forward_context: dict[str, Any],
+) -> None:
+ """
+ Collect Mamba state copy metadata for cross-request copy (e.g., from parent to forked child).
+ This is similar to collect_mamba_copy_meta but uses different request states for source and dest.
+ """
+ src_ptrs_np = copy_bufs.src_ptrs.np
+ dst_ptrs_np = copy_bufs.dst_ptrs.np
+ sizes_np = copy_bufs.sizes.np
+ offset = copy_bufs.offset
+
+ for mamba_group_id in mamba_group_ids:
+ src_block_ids = src_req_state.block_ids[mamba_group_id]
+ dest_block_ids = dest_req_state.block_ids[mamba_group_id]
+ dest_block_id = dest_block_ids[dest_block_idx]
+ layer_names = kv_cache_config.kv_cache_groups[mamba_group_id].layer_names
+ for layer_name in layer_names:
+ attention = forward_context[layer_name]
+ kv_caches: list[torch.Tensor] = attention.kv_cache[0]
+ for state, state_copy_func in zip(kv_caches, mamba_state_copy_funcs):
+ # Get source pointer from parent's block
+ copy_spec = state_copy_func(
+ state, src_block_ids, src_block_idx, 1 # accept_token_bias + 1 = 1 for fresh copy
+ )
+
+ src_ptrs_np[offset] = copy_spec.start_addr
+ dst_ptrs_np[offset] = state[dest_block_id].data_ptr()
+ sizes_np[offset] = copy_spec.num_elements * state.element_size()
+ offset += 1
+
+ copy_bufs.offset = offset
+
+
def postprocess_mamba(
scheduler_output: SchedulerOutput,
kv_cache_config: KVCacheConfig,
diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py
index bede06592f7d..6df8745a500d 100644
--- a/vllm/v1/worker/utils.py
+++ b/vllm/v1/worker/utils.py
@@ -2,7 +2,10 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
from collections import defaultdict
+from collections.abc import Iterable
from dataclasses import dataclass, field
+from itertools import product as iprod
+from typing import Any
import torch
@@ -12,6 +15,8 @@
from vllm.model_executor.models.interfaces import MultiModalEmbeddings
from vllm.model_executor.models.utils import extract_layer_index
from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import largest_power_of_2_divisor
from vllm.utils.mem_utils import MemorySnapshot, format_gib
from vllm.v1.attention.backend import (
AttentionBackend,
@@ -21,6 +26,7 @@
from vllm.v1.kv_cache_interface import (
AttentionSpec,
EncoderOnlyAttentionSpec,
+ FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheSpec,
@@ -31,6 +37,186 @@
logger = init_logger(__name__)
+@triton.jit
+def _zero_kv_blocks_kernel(
+ seg_addrs_ptr,
+ block_ids_ptr,
+ n_blocks,
+ N_SEGS: tl.constexpr,
+ PAGE_SIZE_EL: tl.constexpr,
+ BLOCK_SIZE: tl.constexpr,
+):
+ """Zero KV cache blocks across all segments in a single launch.
+
+ Each segment is a contiguous region of one block's data. For backends
+ where blocks are outermost (block_dim=0) there is one segment per
+ buffer. For backends where K/V is outermost (block_dim=1) there are
+ two segments per buffer (one for K, one for V).
+
+ seg_addrs_ptr holds absolute byte addresses (int64) for each segment,
+ allowing segments to live in different CUDA allocations.
+
+ Programs are mapped as (block_index, seg_index, chunk_index).
+ """
+ pid = tl.program_id(0)
+ chunks = PAGE_SIZE_EL // BLOCK_SIZE
+ work_per_block = N_SEGS * chunks
+ block_index = pid // work_per_block
+ if block_index >= n_blocks:
+ return
+ remainder = pid % work_per_block
+ seg_index = remainder // chunks
+ chunk_index = remainder % chunks
+ block_id = tl.load(block_ids_ptr + block_index)
+ seg_addr = tl.load(seg_addrs_ptr + seg_index)
+ ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32))
+ offset = (
+ block_id.to(tl.int64) * PAGE_SIZE_EL + chunk_index.to(tl.int64) * BLOCK_SIZE
+ )
+ cols = tl.arange(0, BLOCK_SIZE).to(tl.int64)
+ tl.store(ptr + offset + cols, tl.zeros([BLOCK_SIZE], dtype=tl.int32))
+
+
+class KVBlockZeroer:
+ """Manages efficient zeroing of KV cache blocks via a Triton kernel.
+
+ Call :meth:`init_meta` once after KV caches are allocated to precompute
+ segment addresses, then call :meth:`zero_block_ids` each step to zero
+ newly-allocated blocks.
+ """
+
+ def __init__(self, device: torch.device, pin_memory: bool):
+ self.device = device
+ self.pin_memory = pin_memory
+ self._meta: tuple[torch.Tensor, int, int, int] | None = None
+ self._id_cap: int = 0
+ self._ids_pinned: torch.Tensor | None = None
+ self._ids_gpu: torch.Tensor | None = None
+
+ def init_meta(
+ self,
+ attn_groups_iter: Iterable["AttentionGroup"],
+ kernel_block_sizes: list[int],
+ cache_dtype: str,
+ runner_only_attn_layers: set[str],
+ static_forward_context: dict[str, Any],
+ ) -> None:
+ """One-time precomputation for zero_block_ids.
+
+ Builds absolute-address table for the Triton zeroing kernel.
+ Each entry is the absolute byte address of a segment start on the
+ GPU, so segments in different CUDA allocations work correctly.
+
+ Block IDs from the scheduler reference logical blocks whose size
+ may differ from the kernel block size (virtual block splitting).
+ PAGE_SIZE_EL accounts for this ratio so that
+ ``block_id * PAGE_SIZE_EL`` lands at the correct offset.
+
+ Only AttentionSpec layers are processed; Mamba layers are skipped.
+ """
+ seen_ptrs: set[int] = set()
+ seg_addrs: list[int] = []
+ page_size_el: int | None = None
+
+ for group in attn_groups_iter:
+ spec = group.kv_cache_spec
+ if type(spec) is not FullAttentionSpec:
+ continue
+ if group.kv_cache_group_id >= len(kernel_block_sizes):
+ continue
+ kernel_bs = kernel_block_sizes[group.kv_cache_group_id]
+ ratio = spec.block_size // kernel_bs
+ block_dim = group.backend.get_kv_cache_block_dim(
+ kernel_bs,
+ spec.num_kv_heads,
+ spec.head_size,
+ cache_dtype_str=cache_dtype,
+ )
+
+ for layer_name in group.layer_names:
+ if layer_name in runner_only_attn_layers:
+ continue
+ kv = static_forward_context[layer_name].kv_cache[0]
+ if isinstance(kv, list):
+ continue
+ dp = kv.data_ptr()
+ if dp in seen_ptrs:
+ continue
+ seen_ptrs.add(dp)
+
+ el = kv.element_size()
+ cur_bytes = kv.stride(block_dim) * el
+ assert cur_bytes % 4 == 0
+ kernel_block_el = cur_bytes // 4
+ cur_page_el = kernel_block_el * ratio
+ if page_size_el is None:
+ page_size_el = cur_page_el
+ else:
+ assert page_size_el == cur_page_el, (
+ f"Non-uniform page sizes: {page_size_el} vs {cur_page_el}"
+ )
+
+ block_stride_bytes = cur_bytes
+ outer_dims = [
+ d
+ for d in range(block_dim)
+ if kv.stride(d) * el > block_stride_bytes
+ ]
+ outer_strides = [kv.stride(d) * el for d in outer_dims]
+ for outer in iprod(*(range(kv.shape[d]) for d in outer_dims)):
+ off_bytes = sum(i * s for i, s in zip(outer, outer_strides))
+ seg_addrs.append(dp + off_bytes)
+
+ if not seg_addrs or page_size_el is None:
+ self._meta = None
+ return
+
+ blk_size = min(largest_power_of_2_divisor(page_size_el), 1024)
+ self._id_cap = 8192
+ self._ids_pinned = torch.empty(
+ self._id_cap,
+ dtype=torch.int64,
+ pin_memory=self.pin_memory,
+ )
+ self._ids_gpu = torch.empty(self._id_cap, dtype=torch.int64, device=self.device)
+ self._meta = (
+ torch.tensor(seg_addrs, dtype=torch.int64, device=self.device),
+ page_size_el,
+ blk_size,
+ len(seg_addrs),
+ )
+
+ def zero_block_ids(self, block_ids: list[int]) -> None:
+ """Zero the KV cache memory for the given block IDs."""
+ if not block_ids or self._meta is None:
+ return
+ seg_addrs, page_size_el, blk_size, n_segs = self._meta
+ n_blocks = len(block_ids)
+ if n_blocks > self._id_cap:
+ self._id_cap = n_blocks * 2
+ self._ids_pinned = torch.empty(
+ self._id_cap,
+ dtype=torch.int64,
+ pin_memory=self.pin_memory,
+ )
+ self._ids_gpu = torch.empty(
+ self._id_cap, dtype=torch.int64, device=self.device
+ )
+ assert self._ids_pinned is not None and self._ids_gpu is not None
+ self._ids_pinned[:n_blocks].numpy()[:] = block_ids
+ idx = self._ids_gpu[:n_blocks]
+ idx.copy_(self._ids_pinned[:n_blocks], non_blocking=True)
+ grid = (n_blocks * n_segs * (page_size_el // blk_size),)
+ _zero_kv_blocks_kernel[grid](
+ seg_addrs,
+ idx,
+ n_blocks,
+ N_SEGS=n_segs,
+ PAGE_SIZE_EL=page_size_el,
+ BLOCK_SIZE=blk_size,
+ )
+
+
@dataclass
class AttentionGroup:
backend: type[AttentionBackend]